improve loop performance

make loop variable just a pointer, remove make_setter
vyperlang · Jun 8, 2024 · 1485d62 · 1485d62
1 parent 1f6b943
commit 1485d62
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 11 deletions.
diff --git a/vyper/codegen/expr.py b/vyper/codegen/expr.py
@@ -181,6 +181,9 @@ def parse_Name(self):
             ret._referenced_variables = {varinfo}
             return ret
 
+        if varname in self.context.forvars:
+            return self.context.forvars[varname]
+
         if varinfo.is_constant:
             return Expr.parse_value_expr(varinfo.decl_node.value, self.context)
 
@@ -682,7 +685,7 @@ def parse_Call(self):
             (arg0,) = self.expr.args
             arg_ir = Expr(arg0, self.context).ir_node
 
-            assert arg_ir.typ == AddressT()
+            assert isinstance(arg_ir.typ, (AddressT, InterfaceT)), arg_ir.typ
             arg_ir.typ = self.expr._metadata["type"]
 
             return arg_ir

diff --git a/vyper/codegen/stmt.py b/vyper/codegen/stmt.py
@@ -219,13 +219,11 @@ def _parse_For_range(self):
 
             varname = self.stmt.target.target.id
             i = IRnode.from_list(self.context.fresh_varname("range_ix"), typ=target_type)
-            iptr = self.context.new_variable(varname, target_type)
 
-            self.context.forvars[varname] = True
+            assert varname not in forvars
+            self.context.forvars[varname] = i
 
             loop_body = ["seq"]
-            # store the current value of i so it is accessible to userland
-            loop_body.append(["mstore", iptr, i])
             loop_body.append(parse_body(self.stmt.body, self.context))
 
             del self.context.forvars[varname]
@@ -247,14 +245,9 @@ def _parse_For_list(self):
 
         # user-supplied name for loop variable
         varname = self.stmt.target.target.id
-        loop_var = IRnode.from_list(
-            self.context.new_variable(varname, target_type), typ=target_type, location=MEMORY
-        )
 
         i = IRnode.from_list(self.context.fresh_varname("for_list_ix"), typ=UINT256_T)
 
-        self.context.forvars[varname] = True
-
         ret = ["seq"]
 
         # list literal, force it to memory first
@@ -269,7 +262,10 @@ def _parse_For_list(self):
 
         # set up the loop variable
         e = get_element_ptr(iter_list, i, array_bounds_check=False)
-        body = ["seq", make_setter(loop_var, e), parse_body(self.stmt.body, self.context)]
+        assert varname not in self.context.forvars
+        self.context.forvars[varname] = e
+
+        body = parse_body(self.stmt.body, self.context)
 
         repeat_bound = iter_list.typ.count
         if isinstance(iter_list.typ, DArrayT):