fix[lang]: allow flag member access in pure functions (#4693)

flag members are compile-time constants and should be allowed in pure
functions. when validating attribute access like `Action.BUY` in a
pure function, the analysis was failing because `Action` (the value)
is a `TYPE_T` expression which was rejected with an `InvalidReference`
exception.

fix by passing `is_callable=True` to `get_expr_info` when analyzing
the `.value` member of `Attribute` nodes in pure functions. this
allows the `.value` of the attribute (e.g., `Action`) to be parsed
as a type expression without raising `InvalidReference`. note that
this still respects the semantics of `is_callable=True` because we
are only allowing the type expression for the value portion of the
attribute. the overall attribute node (e.g., `Action.BUY`) still
preserves the `is_callable` semantics for the whole expression.

---------

Co-authored-by: Charles Cooper <[email protected]>
Co-authored-by: cyberthirst <[email protected]>

Author: 3 people

tests/functional/codegen/features/test_flag_pure_functions.py

@@ -0,0 +1,150 @@
+def test_flag_members_in_pure_functions(get_contract):
+    """Test that flag members can be used in pure functions since they are
+    compile-time constants"""
+    code = """
+flag Action:
+    BUY
+    SELL
+    CANCEL
+
+@pure
+@external
+def get_buy_action() -> Action:
+    return Action.BUY
+
+@pure
+@external
+def get_sell_action() -> Action:
+    return Action.SELL
+
+@pure
+@external
+def get_cancel_action() -> Action:
+    return Action.CANCEL
+    """
+    c = get_contract(code)
+    assert c.get_buy_action() == 1  # 2^0
+    assert c.get_sell_action() == 2  # 2^1
+    assert c.get_cancel_action() == 4  # 2^2
+
+
+def test_flag_operations_in_pure_functions(get_contract):
+    """Test that flag operations work in pure functions"""
+    code = """
+flag Permissions:
+    READ
+    WRITE
+    EXECUTE
+
+@pure
+@external
+def get_read_write() -> Permissions:
+    return Permissions.READ | Permissions.WRITE
+
+@pure
+@external
+def check_read_permission(perms: Permissions) -> bool:
+    return Permissions.READ in perms
+
+@pure
+@external
+def combine_all() -> Permissions:
+    return Permissions.READ | Permissions.WRITE | Permissions.EXECUTE
+    """
+    c = get_contract(code)
+    assert c.get_read_write() == 3  # 1 | 2 = 3
+    assert c.check_read_permission(1) is True  # READ permission
+    assert c.check_read_permission(2) is False  # WRITE permission only
+    assert c.combine_all() == 7  # 1 | 2 | 4 = 7
+
+
+def test_flag_conditionals_in_pure_functions(get_contract):
+    """Test flags in conditional expressions within pure functions"""
+    code = """
+flag Status:
+    ACTIVE
+    INACTIVE
+    PENDING
+
+@pure
+@external
+def classify_status(status: Status) -> uint256:
+    if status == Status.ACTIVE:
+        return 100
+    elif status == Status.PENDING:
+        return 50
+    else:
+        return 0
+    """
+    c = get_contract(code)
+    assert c.classify_status(1) == 100  # ACTIVE
+    assert c.classify_status(4) == 50  # PENDING
+    assert c.classify_status(2) == 0  # INACTIVE
+
+
+def test_access_flag_from_another_module(get_contract, make_input_bundle):
+    """Test flag access even if the attribute comes from another module (eg lib1.flag.foo)"""
+    code = """
+import lib1
+
+@pure
+@external
+def foo() -> lib1.Action:
+    return lib1.Action.BUY
+
+    """
+    lib1 = """
+flag Action:
+    BUY
+    SELL
+    """
+
+    input_bundle = make_input_bundle({"lib1.vy": lib1})
+    c = get_contract(code, input_bundle=input_bundle)
+    assert c.foo() == 1  # BUY
+
+
+def test_internal_pure_accessing_flag(get_contract, make_input_bundle):
+    """Test flag accesses in internal pure functions"""
+    code = """
+import lib1
+
+@pure
+def bar() -> lib1.Action:
+    return lib1.Action.BUY
+
+@pure
+@external
+def foo() -> lib1.Action:
+    return self.bar()
+
+    """
+    lib1 = """
+flag Action:
+    BUY
+    SELL
+    """
+
+    input_bundle = make_input_bundle({"lib1.vy": lib1})
+    c = get_contract(code, input_bundle=input_bundle)
+    assert c.foo() == 1  # BUY
+
+
+def test_flag_access_in_loop(get_contract, make_input_bundle):
+    """Test flag accesses in a for loop"""
+    code = """
+
+flag Action:
+    BUY
+    SELL
+
+@pure
+@external
+def foo() -> uint256:
+    cnt: uint256 = 0
+    for i: uint256 in range(10):
+        cnt += convert(Action.SELL, uint256)
+    return cnt
+"""
+    c = get_contract(code)
+    assert c.foo() == 10 * 2

tests/functional/codegen/types/test_flag.py

@@ -1,3 +1,8 @@
+import pytest
+
+from vyper.exceptions import StateAccessViolation
+
+
 def test_values_should_be_increasing_ints(get_contract):
     code = """
 flag Action:
@@ -303,3 +308,16 @@ def get_key(f: Foobar, i: uint256) -> uint256:
     """
     c = get_contract(code)
     assert c.get_key(1, 777) == 777
+
+
+@pytest.mark.xfail(raises=StateAccessViolation)
+def test_flag_constant(get_contract):
+    code = """
+flag F:
+    FOO
+    BAR
+
+c: public(constant(F)) = F.FOO
+    """
+    c = get_contract(code)
+    assert c.c() == 1  # FOO

tests/functional/syntax/test_flag.py

@@ -4,6 +4,7 @@
 from vyper.exceptions import (
     FlagDeclarationException,
     InvalidOperation,
+    InvalidReference,
     NamespaceCollision,
     StructureException,
     TypeMismatch,
@@ -124,6 +125,17 @@ def foo():
     """,
         TypeMismatch,
     ),
+    (
+        """
+flag Status:
+  ACTIVE
+
+@external
+def test_assign_to_flag():
+  Status.ACTIVE = 2
+        """,
+        InvalidReference,
+    ),
 ]
 
 

vyper/semantics/analysis/local.py

@@ -179,7 +179,8 @@ def _validate_pure_access(node: vy_ast.Attribute | vy_ast.Name, typ: VyperType)
             raise StateAccessViolation(
                 "not allowed to query environment variables in pure functions"
             )
-        parent_info = get_expr_info(node.value)
+        # allow type exprs in the value node, e.g. MyFlag.A
+        parent_info = get_expr_info(node.value, is_callable=True)
         if isinstance(parent_info.typ, AddressT) and node.attr in AddressT._type_members:
             raise StateAccessViolation("not allowed to query address members in pure functions")
 

vyper/semantics/analysis/utils.py

@@ -90,7 +90,8 @@ def get_expr_info(self, node: vy_ast.VyperNode, is_callable: bool = False) -> Ex
             # propagate the parent exprinfo members down into the new expr
             # note: Attribute(expr value, identifier attr)
 
-            info = self.get_expr_info(node.value, is_callable=is_callable)
+            # allow the value node to be a type expr (e.g., MyFlag.A)
+            info = self.get_expr_info(node.value, is_callable=True)
             attr = node.attr
 
             t = info.typ.get_member(attr, node)