Don't output a lifetime.end for variables - they're valid all the way through a function

src/ast.h

@@ -22,6 +22,7 @@ public:
   SizedType type;
   Map *map = nullptr; // Only set when this expression is assigned to a map
   bool is_literal = false;
+  bool is_variable = false;
 };
 using ExpressionList = std::vector<Expression *>;
 
@@ -71,7 +72,7 @@ public:
 
 class Variable : public Expression {
 public:
-  explicit Variable(std::string &ident) : ident(ident) { }
+  explicit Variable(std::string &ident) : ident(ident) { is_variable = true; }
   std::string ident;
 
   void accept(Visitor &v) override;

src/codegen_llvm.cpp

@@ -257,8 +257,10 @@ void CodegenLLVM::visit(Binop &binop)
       default:
         abort();
     }
-    b_.CreateLifetimeEnd(lhs);
-    b_.CreateLifetimeEnd(rhs);
+    if (!binop.left->is_variable)
+      b_.CreateLifetimeEnd(lhs);
+    if (!binop.right->is_variable)
+      b_.CreateLifetimeEnd(rhs);
   }
   else
   {
@@ -296,12 +298,14 @@ void CodegenLLVM::visit(Unop &unop)
     {
       AllocaInst *dst = b_.CreateAllocaBPF(unop.expr->type, "deref");
       b_.CreateProbeRead(dst, 8, expr_);
-      b_.CreateLifetimeEnd(expr_);
       expr_ = b_.CreateLoad(dst);
+      b_.CreateLifetimeEnd(dst);
       break;
     }
     default: abort();
   }
+
+  // Unops only operate on integer values, so no need for lifetime end
 }
 
 void CodegenLLVM::visit(ExprStatement &expr)
@@ -332,7 +336,8 @@ void CodegenLLVM::visit(AssignMapStatement &assignment)
   }
   b_.CreateMapUpdateElem(map, key, val);
   b_.CreateLifetimeEnd(key);
-  b_.CreateLifetimeEnd(val);
+  if (!assignment.expr->is_variable)
+    b_.CreateLifetimeEnd(val);
 }
 
 void CodegenLLVM::visit(AssignVarStatement &assignment)

tests/codegen.cpp

@@ -1057,6 +1057,88 @@ attributes #1 = { argmemonly nounwind }
 )EXPECTED");
 }
 
+TEST(codegen, variable)
+{
+  test("kprobe:f { $var = comm; @x = $var; @y = $var }",
+
+R"EXPECTED(; Function Attrs: nounwind
+declare i64 @llvm.bpf.pseudo(i64, i64) #0
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.start(i64, i8* nocapture) #1
+
+define i64 @"kprobe:f"(i8* nocapture readnone) local_unnamed_addr section "kprobe:f" {
+entry:
+  %"@y_key" = alloca i64, align 8
+  %"@x_key" = alloca i64, align 8
+  %comm = alloca [64 x i8], align 1
+  %1 = getelementptr inbounds [64 x i8], [64 x i8]* %comm, i64 0, i64 0
+  call void @llvm.lifetime.start(i64 -1, i8* %1)
+  %get_comm = call i64 inttoptr (i64 16 to i64 (i8*, i64)*)([64 x i8]* nonnull %comm, i64 64)
+  %2 = bitcast i64* %"@x_key" to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %2)
+  store i64 0, i64* %"@x_key", align 8
+  %pseudo = call i64 @llvm.bpf.pseudo(i64 1, i64 1)
+  %update_elem = call i64 inttoptr (i64 2 to i64 (i8*, i8*, i8*, i64)*)(i64 %pseudo, i64* nonnull %"@x_key", [64 x i8]* nonnull %comm, i64 0)
+  call void @llvm.lifetime.end(i64 -1, i8* %2)
+  %3 = bitcast i64* %"@y_key" to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %3)
+  store i64 0, i64* %"@y_key", align 8
+  %pseudo1 = call i64 @llvm.bpf.pseudo(i64 1, i64 2)
+  %update_elem2 = call i64 inttoptr (i64 2 to i64 (i8*, i8*, i8*, i64)*)(i64 %pseudo1, i64* nonnull %"@y_key", [64 x i8]* nonnull %comm, i64 0)
+  call void @llvm.lifetime.end(i64 -1, i8* %3)
+  ret i64 0
+}
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.end(i64, i8* nocapture) #1
+
+attributes #0 = { nounwind }
+attributes #1 = { argmemonly nounwind }
+)EXPECTED");
+}
+
+TEST(codegen, dereference)
+{
+  test("kprobe:f { @x = *1234 }",
+
+R"EXPECTED(; Function Attrs: nounwind
+declare i64 @llvm.bpf.pseudo(i64, i64) #0
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.start(i64, i8* nocapture) #1
+
+define i64 @"kprobe:f"(i8* nocapture readnone) local_unnamed_addr section "kprobe:f" {
+entry:
+  %"@x_val" = alloca i64, align 8
+  %"@x_key" = alloca i64, align 8
+  %deref = alloca i64, align 8
+  %1 = bitcast i64* %deref to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %1)
+  %probe_read = call i64 inttoptr (i64 4 to i64 (i8*, i64, i8*)*)(i64* nonnull %deref, i64 8, i64 1234)
+  %2 = load i64, i64* %deref, align 8
+  call void @llvm.lifetime.end(i64 -1, i8* %1)
+  %3 = bitcast i64* %"@x_key" to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %3)
+  store i64 0, i64* %"@x_key", align 8
+  %4 = bitcast i64* %"@x_val" to i8*
+  call void @llvm.lifetime.start(i64 -1, i8* %4)
+  store i64 %2, i64* %"@x_val", align 8
+  %pseudo = call i64 @llvm.bpf.pseudo(i64 1, i64 1)
+  %update_elem = call i64 inttoptr (i64 2 to i64 (i8*, i8*, i8*, i64)*)(i64 %pseudo, i64* nonnull %"@x_key", i64* nonnull %"@x_val", i64 0)
+  call void @llvm.lifetime.end(i64 -1, i8* %3)
+  call void @llvm.lifetime.end(i64 -1, i8* %4)
+  ret i64 0
+}
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.end(i64, i8* nocapture) #1
+
+attributes #0 = { nounwind }
+attributes #1 = { argmemonly nounwind }
+)EXPECTED");
+}
+
 } // namespace codegen
 } // namespace test
 } // namespace bpftrace