Improved trace (#371)

* overriding trace on library defined exprs
algorand · May 31, 2022 · db381c4 · db381c4
1 parent 9646a1a
commit db381c4
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/pyteal/ast/subroutine.py b/pyteal/ast/subroutine.py
@@ -427,4 +427,6 @@ def var_n_loaded(param):
     bodyOps = [var.slot.store() for var in argumentVars[::-1]]
     bodyOps.append(subroutineBody)
 
-    return SubroutineDeclaration(subroutine, Seq(bodyOps))
+    sd = SubroutineDeclaration(subroutine, Seq(bodyOps))
+    sd.trace = subroutineBody.trace
+    return sd
diff --git a/pyteal/compiler/compiler.py b/pyteal/compiler/compiler.py
@@ -128,9 +128,15 @@ def compileSubroutine(
 
     if not ast.has_return():
         if ast.type_of() == TealType.none:
-            ast = Seq([ast, Return()])
+            ret_expr = Return()
+            ret_expr.trace = ast.trace
+            seq_expr = Seq([ast, ret_expr])
+            seq_expr.trace = ret_expr.trace
+            ast = seq_expr
         else:
-            ast = Return(ast)
+            ret_expr = Return(ast)
+            ret_expr.trace = ast.trace
+            ast = ret_expr
 
     options.setSubroutine(currentSubroutine)
     start, end = ast.__teal__(options)