Update logic

lib/Dialect/FIRRTL/Transforms/InferReadWrite.cpp

@@ -210,6 +210,7 @@ struct InferReadWritePass : public InferReadWriteBase<InferReadWritePass> {
             opsToErase.push_back(sf);
           }
       }
+      simplifyWmode(rwMem);
       // All uses for all results of mem removed, now erase the memOp.
       opsToErase.push_back(memOp);
     }
@@ -339,10 +340,6 @@ struct InferReadWritePass : public InferReadWriteBase<InferReadWritePass> {
         }
 
       if (enableDriver && wmodeDriver) {
-        // wmodeDriver must be a wire, get the Value that drives it.
-        wmodeDriver = getConnectSrc(wmodeDriver);
-        if (!wmodeDriver)
-          continue;
         ImplicitLocOpBuilder builder(memOp.getLoc(), memOp);
         auto constOne = builder.create<ConstantOp>(
             UIntType::get(builder.getContext(), 1), APInt(1, 1));
@@ -354,11 +351,21 @@ struct InferReadWritePass : public InferReadWriteBase<InferReadWritePass> {
   // Replace any occurence of enable on the expression tree of wmode with a
   // constant one.
   void setEnable(Value enableDriver, Value wmodeDriver, Value constOne) {
+    auto getDriverOp = [&](Value dst) -> Operation * {
+      // Look through one level of wire to get the driver op.
+      auto *defOp = dst.getDefiningOp();
+      if (defOp) {
+        if (isa<WireOp>(defOp))
+          dst = getConnectSrc(dst);
+        defOp = dst.getDefiningOp();
+      }
+      return defOp;
+    };
     SmallVector<Value> stack;
     stack.push_back(wmodeDriver);
     while (!stack.empty()) {
       auto driver = stack.pop_back_val();
-      auto *defOp = driver.getDefiningOp();
+      auto *defOp = getDriverOp(driver);
       if (!defOp)
         continue;
       for (auto operand : llvm::enumerate(defOp->getOperands())) {