diff --git a/rw_core/flowdecomposition/flow_decomposer.go b/rw_core/flowdecomposition/flow_decomposer.go
index eb338db..80a9604 100644
--- a/rw_core/flowdecomposition/flow_decomposer.go
+++ b/rw_core/flowdecomposition/flow_decomposer.go
@@ -140,13 +140,11 @@
 		logger.Debug("trap-uni")
 
 		//inPortNo is 0 for wildcard input case, do not include upstream port for controller bound flow in input
-		var inPorts []uint32
+		var inPorts = map[uint32]struct{}{inPortNo: {}}
 		if inPortNo == 0 {
 			inPorts = agent.GetWildcardInputPorts(egressHop.Egress) // exclude egress_hop.egress_port.port_no
-		} else {
-			inPorts = []uint32{inPortNo}
 		}
-		for _, inputPort := range inPorts {
+		for inputPort := range inPorts {
 			// Upstream flow on parent (olt) device
 			faParent := &fu.FlowArgs{
 				KV: fu.OfpFlowModArgs{"priority": uint64(flow.Priority), "cookie": flow.Cookie, "meter_id": uint64(meterID), "write_metadata": metadataFromwriteMetadata},
@@ -457,8 +455,11 @@
 		//so that a valid path can be found for the flow
 		nniPorts := agent.GetNNIPorts()
 		if len(nniPorts) > 0 {
-			inPortNo = nniPorts[0]
-			logger.Debugw("assigning-nni-port-as-in-port-for-multicast-flow", log.Fields{"nni": nniPorts[0], "flow:": flow})
+			for port := range nniPorts {
+				inPortNo = port
+				break
+			}
+			logger.Debugw("assigning-nni-port-as-in-port-for-multicast-flow", log.Fields{"nni": inPortNo, "flow:": flow})
 		}
 	}
 	outPortNo := fu.GetOutPort(flow)
diff --git a/rw_core/flowdecomposition/flow_decomposer_test.go b/rw_core/flowdecomposition/flow_decomposer_test.go
index 4c6ca8c..c3bbff7 100644
--- a/rw_core/flowdecomposition/flow_decomposer_test.go
+++ b/rw_core/flowdecomposition/flow_decomposer_test.go
@@ -413,15 +413,11 @@
 	return tfd.defaultRules
 }
 
-func (tfd *testFlowDecomposer) GetWildcardInputPorts(excludePort ...uint32) []uint32 {
-	lPorts := make([]uint32, 0)
-	var exclPort uint32
-	if len(excludePort) == 1 {
-		exclPort = excludePort[0]
-	}
-	for portno := range tfd.logicalPorts {
-		if portno != exclPort {
-			lPorts = append(lPorts, portno)
+func (tfd *testFlowDecomposer) GetWildcardInputPorts(excludePort uint32) map[uint32]struct{} {
+	lPorts := make(map[uint32]struct{})
+	for portNo := range tfd.logicalPorts {
+		if portNo != excludePort {
+			lPorts[portNo] = struct{}{}
 		}
 	}
 	return lPorts
@@ -449,11 +445,11 @@
 	return nil, status.Errorf(codes.FailedPrecondition, "no route from:%d to:%d", ingressPortNo, egressPortNo)
 }
 
-func (tfd *testFlowDecomposer) GetNNIPorts() []uint32 {
-	nniPorts := make([]uint32, 0)
+func (tfd *testFlowDecomposer) GetNNIPorts() map[uint32]struct{} {
+	nniPorts := make(map[uint32]struct{})
 	for portNo, nni := range tfd.logicalPortsNo {
 		if nni {
-			nniPorts = append(nniPorts, portNo)
+			nniPorts[portNo] = struct{}{}
 		}
 	}
 	return nniPorts