gollvm: fix bug in call->invoke instruction rewriting

During the control flow generation pass in gollvm, there is a pass
that rewrites call instructions to invoke instructions in situations
where this is needed due to the presence of a defer. The rewriter
wasn't properly handling the case where a given Bexpression was
inheriting the LLVM value of a child expression, and as a result the
compiler was crashing when trying to access a stale (deleted) LLVM
value.

Fixes golang/go#31388.

Change-Id: I6abc5ec87b3f63d2974126355a9d8fe18ec5ed1f
Reviewed-on: https://go-review.googlesource.com/c/gollvm/+/171357
Reviewed-by: Cherry Zhang <cherryyz@google.com>
diff --git a/bridge/go-llvm-bnode.cpp b/bridge/go-llvm-bnode.cpp
index a3a1ac9..0e8e006 100644
--- a/bridge/go-llvm-bnode.cpp
+++ b/bridge/go-llvm-bnode.cpp
@@ -1059,3 +1059,16 @@
     }
   }
 }
+
+void BnodeBuilder::updateValue(Bexpression *expr, llvm::Value *newValue)
+{
+  assert(expr != NULL && newValue != NULL && expr->value() != NULL);
+
+  // New value should be different from old value
+  assert(expr->value() != newValue);
+
+  // Types should match
+  assert(expr->value()->getType() == newValue->getType());
+
+  expr->setValue(newValue);
+}
diff --git a/bridge/go-llvm-bnode.h b/bridge/go-llvm-bnode.h
index ca63b7f..fff709f 100644
--- a/bridge/go-llvm-bnode.h
+++ b/bridge/go-llvm-bnode.h
@@ -381,6 +381,13 @@
   void updateInstructions(Bexpression *expr,
                           std::vector<llvm::Instruction*> newinsts);
 
+
+  // Similar to the above, but used in cases where a parent Bexpression
+  // inherits the value from a child, and the child's value is rewritten
+  // to some new value (simpler case here, since there is no book-keeping
+  // update needed for the integrity checker).
+  void updateValue(Bexpression *expr, llvm::Value *newValue);
+
   // Get/set whether the tree integrity checker is enabled. It makes sense
   // to turn off the integrity checker during tree cloning operations
   // (part of sharing repair), and also for unit testing.
diff --git a/bridge/go-llvm.cpp b/bridge/go-llvm.cpp
index 53bccbe..86e56e8 100644
--- a/bridge/go-llvm.cpp
+++ b/bridge/go-llvm.cpp
@@ -2772,7 +2772,8 @@
   llvm::BasicBlock *getBlockForLabel(Blabel *lab);
   llvm::BasicBlock *walkExpr(llvm::BasicBlock *curblock,
                              Bstatement *containingStmt,
-                             Bexpression *expr);
+                             Bexpression *expr,
+                             bool subexpr=false);
   std::pair<llvm::Instruction*, llvm::BasicBlock *>
   rewriteToMayThrowCall(llvm::CallInst *call,
                         llvm::BasicBlock *curblock);
@@ -2809,6 +2810,7 @@
   std::vector<llvm::BasicBlock*> padBlockStack_;
   std::set<llvm::Instruction *> temporariesDiscovered_;
   std::vector<llvm::Instruction *> newTemporaries_;
+  std::map<llvm::Value *, llvm::Instruction *> instRewrites_;
   llvm::BasicBlock *finallyBlock_;
   Bstatement *cachedReturn_;
 };
@@ -2992,12 +2994,36 @@
   // return value.
   call->replaceAllUsesWith(invcall);
 
+  // Add an entry in the rewrites map to record the fact that we've replaced the
+  // call with an invoke. This is to take care of situations where the LLVM
+  // value for a given Bexpression is inherited by the expression's parent.
+  // Example:
+  //
+  //    func f(...) int64 { ... }
+  //    ...
+  //    func g() {
+  //      defer func() { ... }()
+  //      z = uint64(f(...))
+  //      ...
+  //    }
+  //
+  // This will result in a call parented by a conversion, e.g.
+  //
+  //    conv: [ btype: [uns] 'uint64' i64 ]
+  //       call: [ btype: 'int64' i64 ]
+  //          fcn: [ btype: i64 (i8*, i64)* ] main.F
+  //          ...
+  //
+  // Both the call and the convert will wind up with the same LLVM Value,
+  // hence we need to rewrite the value not just of the Bexpression for the
+  // call, but also the Bexpression for the convert.
+  //
+  assert(instRewrites_.find(call) == instRewrites_.end());
+  instRewrites_[call] = invcall;
+
   // New call needs same attributes
   invcall->setAttributes(call->getAttributes());
 
-  // Old call longer needed
-  call->deleteValue();
-
   return std::make_pair(invcall, contbb);
 }
 
@@ -3053,7 +3079,8 @@
 
 llvm::BasicBlock *GenBlocks::walkExpr(llvm::BasicBlock *curblock,
                                       Bstatement *containingStmt,
-                                      Bexpression *expr)
+                                      Bexpression *expr,
+                                      bool subexpr)
 {
   // Delete dead instructions before visiting the children,
   // as they may use values defined in the children. Uses
@@ -3063,8 +3090,13 @@
 
   // Visit children first
   const std::vector<Bnode *> &kids = expr->children();
-  for (auto &child : kids)
-    curblock = walk(child, containingStmt, curblock);
+  for (auto &child : kids) {
+    Bexpression *expr = child->castToBexpression();
+    if (expr)
+      curblock = walkExpr(curblock, containingStmt, expr, true);
+    else
+      curblock = walk(child, containingStmt, curblock);
+  }
 
   // In case it becomes dead after visiting some child...
   if (!curblock)
@@ -3101,6 +3133,22 @@
   if (changed)
     be_->nodeBuilder().updateInstructions(expr, newinsts);
   expr->clear();
+
+  // Rewrite value for this expr if needed.
+  auto it = instRewrites_.find(expr->value());
+  if (it != instRewrites_.end())
+    be_->nodeBuilder().updateValue(expr, it->second);
+
+  // If this is a top-level expression (e.g. expression parented by statement,
+  // as opposed to expression parented by some other expr), then at this
+  // point we can delete any inst that appears as a key in the rewrite
+  // map (we don't expect to see any other uses).
+  if (!subexpr) {
+    for (auto it = instRewrites_.begin(); it != instRewrites_.end(); it++)
+      it->first->deleteValue();
+    instRewrites_.clear();
+  }
+
   return curblock;
 }