From 9e0c1696f3caef4059c65774ad6b8efee91d9d9e Mon Sep 17 00:00:00 2001
From: Andrew Lenharth <andrew@lenharth.org>
Date: Fri, 16 Feb 2024 10:45:51 -0600
Subject: [PATCH] [HW] Fix more bugs related to input/output specific
 functions.

Fix more hold-overs from when input and output were split.

See #6707
---
 include/circt/Dialect/HW/HWOpInterfaces.td | 36 +++++++++++++---------
 lib/Dialect/ESI/ESIOps.cpp                 |  2 +-
 lib/Dialect/HW/HWOps.cpp                   | 18 ++++-------
 3 files changed, 28 insertions(+), 28 deletions(-)

diff --git a/include/circt/Dialect/HW/HWOpInterfaces.td b/include/circt/Dialect/HW/HWOpInterfaces.td
index 422400fec1..7d3bd44f9e 100644
--- a/include/circt/Dialect/HW/HWOpInterfaces.td
+++ b/include/circt/Dialect/HW/HWOpInterfaces.td
@@ -65,7 +65,7 @@ def HWModuleLike : OpInterface<"HWModuleLike", [
     "SmallVector<Location>", "getAllPortLocs", (ins)>,
 
     InterfaceMethod<"Set the port Locations",
-    "void", "setAllPortLocs", (ins "ArrayRef<Location>":$locs)>,
+    "void", "setAllPortLocsAttrs", (ins "ArrayRef<Attribute>":$locs)>,
 
     InterfaceMethod<"Set the module type (and port names)",
     "void", "setHWModuleType", (ins "::circt::hw::ModuleType":$type)>,
@@ -259,6 +259,11 @@ def HWModuleLike : OpInterface<"HWModuleLike", [
       $_op.setAllPortAttrs(attrs);
     }
 
+    void setAllPortLocs(ArrayRef<Location> locs) {
+      std::vector<Attribute> nLocs(locs.begin(), locs.end());
+      $_op.setAllPortLocsAttrs(nLocs);
+    }
+
     Location getPortLoc(size_t idx) {
       return $_op.getAllPortLocs()[idx];
     }
@@ -293,13 +298,14 @@ def HWModuleLike : OpInterface<"HWModuleLike", [
       return ArrayAttr::get($_op->getContext(), retval);
     }
 
-    void setInputLocs(ArrayRef<Location> inAttrs) {
-      assert(inAttrs.size() == $_op.getNumInputPorts());
-      auto outAttrs = getOutputLocs();
-      SmallVector<Location> attrs;
-      attrs.append(inAttrs.begin(), inAttrs.end());
-      attrs.append(outAttrs.begin(), outAttrs.end());
-      $_op.setAllPortLocs(attrs);
+    void setInputLocs(ArrayRef<Location> inLocs) {
+      assert(inLocs.size() == $_op.getNumInputPorts());
+      auto type = $_op.getHWModuleType();
+      SmallVector<Location> newLocs($_op.getAllPortLocs());
+      for (size_t idx = 0, e = inLocs.size(); idx != e; ++idx)
+        newLocs[type.getPortIdForInputId(idx)] = inLocs[idx];
+      $_op.setAllPortLocs(newLocs);
+
     }
 
     SmallVector<Location> getOutputLocs() {
@@ -318,13 +324,13 @@ def HWModuleLike : OpInterface<"HWModuleLike", [
       return ArrayAttr::get($_op->getContext(), retval);
     }
 
-    void setOutputLocs(ArrayRef<Location> outAttrs) {
-      assert(outAttrs.size() == $_op.getNumOutputPorts());
-      auto inAttrs = getInputLocs();
-      SmallVector<Location> attrs;
-      attrs.append(inAttrs.begin(), inAttrs.end());
-      attrs.append(outAttrs.begin(), outAttrs.end());
-      $_op.setAllPortLocs(attrs);
+    void setOutputLocs(ArrayRef<Location> outLocs) {
+      assert(outLocs.size() == $_op.getNumOutputPorts());
+      auto type = $_op.getHWModuleType();
+      SmallVector<Location> newLocs($_op.getAllPortLocs());
+      for (size_t idx = 0, e = outLocs.size(); idx != e; ++idx)
+        newLocs[type.getPortIdForOutputId(idx)] = outLocs[idx];
+      $_op.setAllPortLocs(newLocs);
     }
   }];
 
diff --git a/lib/Dialect/ESI/ESIOps.cpp b/lib/Dialect/ESI/ESIOps.cpp
index dcbbfbe663..f3e42ea3df 100644
--- a/lib/Dialect/ESI/ESIOps.cpp
+++ b/lib/Dialect/ESI/ESIOps.cpp
@@ -598,7 +598,7 @@ SmallVector<Location> ESIPureModuleOp::getAllPortLocs() {
   return retval;
 }
 
-void ESIPureModuleOp::setAllPortLocs(ArrayRef<Location> locs) {
+void ESIPureModuleOp::setAllPortLocsAttrs(ArrayRef<Attribute> locs) {
   emitError("No ports for port locations");
 }
 
diff --git a/lib/Dialect/HW/HWOps.cpp b/lib/Dialect/HW/HWOps.cpp
index f0a70bfe04..60427b9a60 100644
--- a/lib/Dialect/HW/HWOps.cpp
+++ b/lib/Dialect/HW/HWOps.cpp
@@ -1246,13 +1246,7 @@ SmallVector<Location> HWModuleGeneratedOp::getAllPortLocs() {
   return ::getAllPortLocs(*this);
 }
 
-template <typename ModTy>
-static void setAllPortLocs(ArrayRef<Location> locs, ModTy module) {
-  std::vector<Attribute> nLocs(locs.begin(), locs.end());
-  module.setPortLocsAttr(ArrayAttr::get(module.getContext(), nLocs));
-}
-
-void HWModuleOp::setAllPortLocs(ArrayRef<Location> locs) {
+void HWModuleOp::setAllPortLocsAttrs(ArrayRef<Attribute> locs) {
   SmallVector<Attribute> resultLocs;
   unsigned inputCount = 0;
   auto modType = getModuleType();
@@ -1261,17 +1255,17 @@ void HWModuleOp::setAllPortLocs(ArrayRef<Location> locs) {
     if (modType.isOutput(i))
       resultLocs.push_back(locs[i]);
     else
-      body->getArgument(inputCount++).setLoc(locs[i]);
+      body->getArgument(inputCount++).setLoc(cast<Location>(locs[i]));
   }
   setResultLocsAttr(ArrayAttr::get(getContext(), resultLocs));
 }
 
-void HWModuleExternOp::setAllPortLocs(ArrayRef<Location> locs) {
-  ::setAllPortLocs(locs, *this);
+void HWModuleExternOp::setAllPortLocsAttrs(ArrayRef<Attribute> locs) {
+  setPortLocsAttr(ArrayAttr::get(getContext(), locs));
 }
 
-void HWModuleGeneratedOp::setAllPortLocs(ArrayRef<Location> locs) {
-  ::setAllPortLocs(locs, *this);
+void HWModuleGeneratedOp::setAllPortLocsAttrs(ArrayRef<Attribute> locs) {
+  setPortLocsAttr(ArrayAttr::get(getContext(), locs));
 }
 
 template <typename ModTy>