cinderx 2026.1.16.2__cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

cinderx/compiler/static/effects.py

@@ -0,0 +1,160 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# pyre-strict
+
+from __future__ import annotations
+
+import ast
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .types import Value
+
+    RefinedFields = dict[str, tuple[Value, int, set[ast.AST]]]
+
+
+# A refined field consists of a refined type in addition to the node that
+# refined the field. The node information is used to hoist reads during codegen.
+
+
+class TypeState:
+    def __init__(self) -> None:
+        self.local_types: dict[str, Value] = {}
+        self.refined_fields: dict[str, RefinedFields] = {}
+
+    def copy(self) -> TypeState:
+        type_state = TypeState()
+        type_state.local_types = dict(self.local_types)
+        type_state.refined_fields = dict(self.refined_fields)
+        return type_state
+
+    def __repr__(self) -> str:
+        return f"TypeState({self.local_types}, {self.refined_fields})"
+
+
+class NarrowingEffect:
+    """captures type narrowing effects on variables"""
+
+    def and_(self, other: NarrowingEffect) -> NarrowingEffect:
+        if other is NO_EFFECT:
+            return self
+
+        return AndEffect(self, other)
+
+    def or_(self, other: NarrowingEffect) -> NarrowingEffect:
+        if other is NO_EFFECT:
+            return self
+
+        return OrEffect(self, other)
+
+    def not_(self) -> NarrowingEffect:
+        return NegationEffect(self)
+
+    def apply(
+        self,
+        type_state: TypeState,
+        type_state_nodes: dict[str, ast.AST] | None = None,
+    ) -> None:
+        """applies the given effect in the target scope. if `type_state_nodes` is passed, populates
+        it with the underlying name or attribute nodes"""
+        pass
+
+    def undo(self, type_state: TypeState) -> None:
+        """restores the type to its original value"""
+        pass
+
+    def reverse(
+        self,
+        type_state: TypeState,
+        type_state_nodes: dict[str, ast.AST] | None = None,
+    ) -> None:
+        """applies the reverse of the scope or reverts it if
+        there is no reverse"""
+        self.undo(type_state)
+
+
+class AndEffect(NarrowingEffect):
+    def __init__(self, *effects: NarrowingEffect) -> None:
+        self.effects: Sequence[NarrowingEffect] = effects
+
+    def and_(self, other: NarrowingEffect) -> NarrowingEffect:
+        if other is NO_EFFECT:
+            return self
+        elif isinstance(other, AndEffect):
+            return AndEffect(*self.effects, *other.effects)
+
+        return AndEffect(*self.effects, other)
+
+    def apply(
+        self,
+        type_state: TypeState,
+        type_state_nodes: dict[str, ast.AST] | None = None,
+    ) -> None:
+        for effect in self.effects:
+            effect.apply(type_state, type_state_nodes)
+
+    def undo(self, type_state: TypeState) -> None:
+        """restores the type to its original value"""
+        for effect in self.effects:
+            effect.undo(type_state)
+
+
+class OrEffect(NarrowingEffect):
+    def __init__(self, *effects: NarrowingEffect) -> None:
+        self.effects: Sequence[NarrowingEffect] = effects
+
+    def and_(self, other: NarrowingEffect) -> NarrowingEffect:
+        if other is NO_EFFECT:
+            return self
+        elif isinstance(other, OrEffect):
+            return OrEffect(*self.effects, *other.effects)
+
+        return OrEffect(*self.effects, other)
+
+    def reverse(
+        self,
+        type_state: TypeState,
+        type_state_nodes: dict[str, ast.AST] | None = None,
+    ) -> None:
+        for effect in self.effects:
+            effect.reverse(type_state, type_state_nodes)
+
+    def undo(self, type_state: TypeState) -> None:
+        """restores the type to its original value"""
+        for effect in self.effects:
+            effect.undo(type_state)
+
+
+class NoEffect(NarrowingEffect):
+    def union(self, other: NarrowingEffect) -> NarrowingEffect:
+        return other
+
+
+# Singleton instance for no effects
+NO_EFFECT = NoEffect()
+
+
+class NegationEffect(NarrowingEffect):
+    def __init__(self, negated: NarrowingEffect) -> None:
+        self.negated = negated
+
+    def not_(self) -> NarrowingEffect:
+        return self.negated
+
+    def apply(
+        self,
+        type_state: TypeState,
+        type_state_nodes: dict[str, ast.AST] | None = None,
+    ) -> None:
+        self.negated.reverse(type_state, type_state_nodes)
+
+    def undo(self, type_state: TypeState) -> None:
+        self.negated.undo(type_state)
+
+    def reverse(
+        self,
+        type_state: TypeState,
+        type_state_nodes: dict[str, ast.AST] | None = None,
+    ) -> None:
+        self.negated.apply(type_state, type_state_nodes)