Fix segfault when lowering ConstructOp to ZHLT

zirgen/Conversions/Typing/ZhlComponent.cpp

@@ -895,9 +895,27 @@ void LoweringImpl::gen(ConstructOp construct, ComponentBuilder& cb) {
     return;
   }
 
-  Value layout =
-      addOrExpandLayoutMember(construct.getLoc(), cb, construct.getOut(), ctor.getLayoutType());
-  auto call = builder.create<Zhlt::ConstructOp>(construct.getLoc(), ctor, arguments, layout);
+  // Don't expand the member layout until the actual definition -- even in the
+  // arguments to the ConstructOp. This ensures that we don't leave any dangling
+  // uses of the "old" layout after we expand it, since the declaration should
+  // not be referenced after the definition.
+  Zhlt::ConstructOp call;
+  Value tmpLayout;
+  if (ctor.getLayoutType()) {
+      tmpLayout = builder.create<Zhlt::MagicOp>(construct.getLoc(), ctor.getLayoutType());
+  }
+  call = builder.create<Zhlt::ConstructOp>(construct.getLoc(), ctor, arguments, tmpLayout);
+
+  // Now that we've built the constructor call, expand the layout and clean up
+  // the temporary one if necessary.
+  Value layout;
+  if (tmpLayout) {
+    layout =
+        addOrExpandLayoutMember(construct.getLoc(), cb, construct.getOut(), ctor.getLayoutType());
+    tmpLayout.replaceAllUsesWith(layout);
+    tmpLayout.getDefiningOp()->erase();
+  }
+
   if (layout && ctor->getAttr("alwaysinline")) {
     construct.emitError() << "Cannot construct non-trivial layouts inside of a function";
     throw MalformedIRException();

zirgen/dsl/test/parameter_backs.zir

@@ -0,0 +1,26 @@
+// RUN: zirgen %s --test --test-cycles=4 | FileCheck %s
+
+// CHECK: [0] Log: result = 1
+// CHECK: [1] Log: result = 2
+// CHECK: [2] Log: result = 4
+// CHECK: [3] Log: result = 8
+
+extern IsFirstCycle() : Val;
+
+component DoubleBackOne(x: NondetReg) {
+    NondetReg(2 * x@1)
+}
+
+component Top() {
+    first := NondetReg(IsFirstCycle());
+    result : NondetReg;
+    result := [first, 1 - first] -> (
+        NondetReg(1),
+        DoubleBackOne(result@0)
+    );
+}
+
+test {
+    top := Top();
+    Log("result = %u", top.result);
+}

zirgen/dsl/test/repro/back_arg_segfault.zir

@@ -0,0 +1,10 @@
+// RUN: zirgen %s --emit=zhlt
+
+component Foo(rt1: Reg) {
+    Reg(5)
+}
+
+component Top() {
+    x : Reg;
+    x := Foo(x@1);
+}