yuho 5.0.0__py3-none-any.whl

yuho/ast/constant_folder.py

@@ -0,0 +1,280 @@
+"""
+Constant folding optimization for Yuho AST.
+
+Evaluates compile-time constant expressions and replaces them with
+their literal values. This includes:
+- Arithmetic operations on numeric literals (+, -, *, /)
+- Boolean operations on boolean literals (&&, ||, !)
+- Comparison operations on numeric literals (==, !=, <, >, <=, >=)
+- String concatenation on string literals (+)
+- Negation of numeric literals (-)
+"""
+
+from typing import Optional, Union
+
+from yuho.ast import nodes
+from yuho.ast.transformer import Transformer
+
+
+class ConstantFoldingError(Exception):
+    """Error during constant folding (e.g., division by zero)."""
+    pass
+
+
+class ConstantFolder(Transformer):
+    """
+    AST transformer that evaluates compile-time constant expressions.
+
+    Traverses the AST and replaces constant expressions with their
+    computed literal values. Preserves source location information
+    for error reporting.
+
+    Usage:
+        folder = ConstantFolder()
+        optimized_ast = folder.transform(ast)
+
+    Handles:
+        - Integer arithmetic: 1 + 2 -> 3
+        - Float arithmetic: 1.5 * 2.0 -> 3.0
+        - Boolean logic: TRUE && FALSE -> FALSE
+        - Comparisons: 5 > 3 -> TRUE
+        - String concatenation: "a" + "b" -> "ab"
+        - Unary negation: -5 -> IntLit(-5)
+        - Unary not: !TRUE -> FALSE
+    """
+
+    def __init__(self, *, fold_division_by_zero: bool = False):
+        """
+        Initialize the constant folder.
+
+        Args:
+            fold_division_by_zero: If True, division by zero will raise
+                ConstantFoldingError. If False (default), the expression
+                is left unchanged for runtime error handling.
+        """
+        self.fold_division_by_zero = fold_division_by_zero
+
+    def transform_binary_expr(self, node: nodes.BinaryExprNode) -> nodes.ASTNode:
+        """
+        Transform binary expressions, folding constants where possible.
+
+        First transforms children (bottom-up folding), then attempts
+        to evaluate if both operands are literals.
+        """
+        # First, transform children to fold nested expressions
+        node = super().transform_binary_expr(node)
+
+        # Check if both operands are now literals
+        left = node.left
+        right = node.right
+        op = node.operator
+
+        # Integer operations
+        if isinstance(left, nodes.IntLit) and isinstance(right, nodes.IntLit):
+            result = self._fold_int_binary(left.value, op, right.value, node)
+            if result is not None:
+                return result
+
+        # Float operations (including mixed int/float)
+        if self._is_numeric_literal(left) and self._is_numeric_literal(right):
+            left_val = self._get_numeric_value(left)
+            right_val = self._get_numeric_value(right)
+            result = self._fold_float_binary(left_val, op, right_val, node)
+            if result is not None:
+                return result
+
+        # Boolean operations
+        if isinstance(left, nodes.BoolLit) and isinstance(right, nodes.BoolLit):
+            result = self._fold_bool_binary(left.value, op, right.value, node)
+            if result is not None:
+                return result
+
+        # String concatenation
+        if isinstance(left, nodes.StringLit) and isinstance(right, nodes.StringLit):
+            if op == "+":
+                return nodes.StringLit(
+                    value=left.value + right.value,
+                    source_location=node.source_location,
+                )
+
+        # Cannot fold, return transformed node
+        return node
+
+    def transform_unary_expr(self, node: nodes.UnaryExprNode) -> nodes.ASTNode:
+        """
+        Transform unary expressions, folding constants where possible.
+
+        Handles:
+            - Negation of numeric literals: -5 -> IntLit(-5)
+            - Logical not of booleans: !TRUE -> FALSE
+        """
+        # First, transform the operand
+        node = super().transform_unary_expr(node)
+
+        operand = node.operand
+        op = node.operator
+
+        # Numeric negation
+        if op == "-":
+            if isinstance(operand, nodes.IntLit):
+                return nodes.IntLit(
+                    value=-operand.value,
+                    source_location=node.source_location,
+                )
+            if isinstance(operand, nodes.FloatLit):
+                return nodes.FloatLit(
+                    value=-operand.value,
+                    source_location=node.source_location,
+                )
+
+        # Logical not
+        if op == "!":
+            if isinstance(operand, nodes.BoolLit):
+                return nodes.BoolLit(
+                    value=not operand.value,
+                    source_location=node.source_location,
+                )
+
+        # Cannot fold, return transformed node
+        return node
+
+    def _is_numeric_literal(self, node: nodes.ASTNode) -> bool:
+        """Check if a node is a numeric literal (int or float)."""
+        return isinstance(node, (nodes.IntLit, nodes.FloatLit))
+
+    def _get_numeric_value(self, node: nodes.ASTNode) -> Union[int, float]:
+        """Extract numeric value from an IntLit or FloatLit."""
+        if isinstance(node, nodes.IntLit):
+            return node.value
+        if isinstance(node, nodes.FloatLit):
+            return node.value
+        raise TypeError(f"Expected numeric literal, got {type(node).__name__}")
+
+    def _fold_int_binary(
+        self,
+        left: int,
+        op: str,
+        right: int,
+        node: nodes.BinaryExprNode,
+    ) -> Optional[nodes.ASTNode]:
+        """Fold binary operations on two integers."""
+        loc = node.source_location
+
+        # Arithmetic operations
+        if op == "+":
+            return nodes.IntLit(value=left + right, source_location=loc)
+        if op == "-":
+            return nodes.IntLit(value=left - right, source_location=loc)
+        if op == "*":
+            return nodes.IntLit(value=left * right, source_location=loc)
+        if op == "/":
+            if right == 0:
+                if self.fold_division_by_zero:
+                    raise ConstantFoldingError(
+                        f"Division by zero at {loc}" if loc else "Division by zero"
+                    )
+                return None  # Leave for runtime
+            # Integer division
+            return nodes.IntLit(value=left // right, source_location=loc)
+
+        # Comparison operations
+        if op == "==":
+            return nodes.BoolLit(value=left == right, source_location=loc)
+        if op == "!=":
+            return nodes.BoolLit(value=left != right, source_location=loc)
+        if op == "<":
+            return nodes.BoolLit(value=left < right, source_location=loc)
+        if op == ">":
+            return nodes.BoolLit(value=left > right, source_location=loc)
+        if op == "<=":
+            return nodes.BoolLit(value=left <= right, source_location=loc)
+        if op == ">=":
+            return nodes.BoolLit(value=left >= right, source_location=loc)
+
+        return None
+
+    def _fold_float_binary(
+        self,
+        left: Union[int, float],
+        op: str,
+        right: Union[int, float],
+        node: nodes.BinaryExprNode,
+    ) -> Optional[nodes.ASTNode]:
+        """Fold binary operations involving at least one float."""
+        loc = node.source_location
+
+        # If both are actually ints, let _fold_int_binary handle it
+        if isinstance(left, int) and isinstance(right, int):
+            return None
+
+        # Arithmetic operations (result is float)
+        if op == "+":
+            return nodes.FloatLit(value=float(left + right), source_location=loc)
+        if op == "-":
+            return nodes.FloatLit(value=float(left - right), source_location=loc)
+        if op == "*":
+            return nodes.FloatLit(value=float(left * right), source_location=loc)
+        if op == "/":
+            if right == 0:
+                if self.fold_division_by_zero:
+                    raise ConstantFoldingError(
+                        f"Division by zero at {loc}" if loc else "Division by zero"
+                    )
+                return None  # Leave for runtime
+            return nodes.FloatLit(value=float(left) / float(right), source_location=loc)
+
+        # Comparison operations (result is bool)
+        if op == "==":
+            return nodes.BoolLit(value=left == right, source_location=loc)
+        if op == "!=":
+            return nodes.BoolLit(value=left != right, source_location=loc)
+        if op == "<":
+            return nodes.BoolLit(value=left < right, source_location=loc)
+        if op == ">":
+            return nodes.BoolLit(value=left > right, source_location=loc)
+        if op == "<=":
+            return nodes.BoolLit(value=left <= right, source_location=loc)
+        if op == ">=":
+            return nodes.BoolLit(value=left >= right, source_location=loc)
+
+        return None
+
+    def _fold_bool_binary(
+        self,
+        left: bool,
+        op: str,
+        right: bool,
+        node: nodes.BinaryExprNode,
+    ) -> Optional[nodes.ASTNode]:
+        """Fold binary operations on two booleans."""
+        loc = node.source_location
+
+        if op == "&&":
+            return nodes.BoolLit(value=left and right, source_location=loc)
+        if op == "||":
+            return nodes.BoolLit(value=left or right, source_location=loc)
+        if op == "==":
+            return nodes.BoolLit(value=left == right, source_location=loc)
+        if op == "!=":
+            return nodes.BoolLit(value=left != right, source_location=loc)
+
+        return None
+
+
+def fold_constants(ast: nodes.ASTNode, **kwargs) -> nodes.ASTNode:
+    """
+    Convenience function to fold constants in an AST.
+
+    Args:
+        ast: The AST to transform.
+        **kwargs: Options passed to ConstantFolder constructor.
+
+    Returns:
+        A new AST with constant expressions folded.
+
+    Example:
+        from yuho.ast.constant_folder import fold_constants
+        optimized = fold_constants(parsed_ast)
+    """
+    folder = ConstantFolder(**kwargs)
+    return folder.transform(ast)

yuho/ast/dead_code.py

@@ -0,0 +1,199 @@
+"""
+Dead code elimination for Yuho AST.
+
+Removes provably unreachable code branches:
+- Match arms that can never be reached (covered by earlier patterns)
+- Branches guarded by constant FALSE conditions
+- Entire match expressions that reduce to a single arm
+"""
+
+from dataclasses import dataclass
+from typing import List, Optional, Set
+
+from yuho.ast import nodes
+from yuho.ast.transformer import Transformer
+from yuho.ast.constant_folder import ConstantFolder
+from yuho.ast.reachability import ReachabilityChecker, ReachabilityResult
+from yuho.ast.type_inference import TypeInferenceResult
+
+
+@dataclass
+class EliminationStats:
+    """Statistics about eliminated dead code."""
+
+    removed_match_arms: int = 0
+    removed_true_guards: int = 0
+    simplified_matches: int = 0
+
+    @property
+    def total_eliminations(self) -> int:
+        return self.removed_match_arms + self.removed_true_guards + self.simplified_matches
+
+
+class DeadCodeEliminator(Transformer):
+    """
+    AST transformer that removes provably unreachable branches.
+
+    Performs the following optimizations:
+    1. Removes match arms that are unreachable (covered by earlier patterns)
+    2. Removes guards that are constant TRUE (simplifies to unguarded arm)
+    3. Removes entire match expressions that have only one reachable arm
+    4. Removes branches guarded by FALSE conditions (not yet if-exprs exist)
+
+    Usage:
+        eliminator = DeadCodeEliminator()
+        optimized_ast = eliminator.transform(ast)
+        print(f"Removed {eliminator.stats.total_eliminations} dead code items")
+
+    Note:
+        This transformer should typically be run after constant folding
+        to maximize dead code detection.
+    """
+
+    def __init__(
+        self,
+        type_info: Optional[TypeInferenceResult] = None,
+        *,
+        fold_constants_first: bool = True,
+    ):
+        """
+        Initialize the dead code eliminator.
+
+        Args:
+            type_info: Type inference result for better reachability analysis.
+            fold_constants_first: If True, fold constants before elimination
+                to maximize dead code detection.
+        """
+        self.type_info = type_info
+        self.fold_constants_first = fold_constants_first
+        self.stats = EliminationStats()
+        self._constant_folder = ConstantFolder() if fold_constants_first else None
+
+    def transform(self, node: nodes.ASTNode) -> nodes.ASTNode:
+        """
+        Transform the AST, eliminating dead code.
+
+        If fold_constants_first is True, runs constant folding pass first.
+        """
+        if self._constant_folder:
+            node = self._constant_folder.transform(node)
+        return super().transform(node)
+
+    def transform_match_expr(self, node: nodes.MatchExprNode) -> nodes.ASTNode:
+        """
+        Transform match expression, removing unreachable arms.
+
+        1. Transform children first (recurse into arm bodies)
+        2. Run reachability analysis
+        3. Filter out unreachable arms
+        4. If only one arm remains and it's a wildcard, simplify
+        """
+        # First, transform children (arm bodies, guards, etc.)
+        node = super().transform_match_expr(node)
+
+        # Run reachability check on this match
+        checker = ReachabilityChecker(self.type_info)
+        result = checker._check_match_reachability(node)
+
+        # Get indices of unreachable arms
+        unreachable_indices: Set[int] = {
+            arm.arm_index for arm in result.unreachable_arms
+        }
+
+        # Filter and transform arms
+        new_arms = []
+        for i, arm in enumerate(node.arms):
+            if i in unreachable_indices:
+                self.stats.removed_match_arms += 1
+                continue
+
+            # Check if guard is constant TRUE - can remove guard
+            new_arm = self._simplify_arm_guard(arm)
+            new_arms.append(new_arm)
+
+        # If no arms remain (shouldn't happen in valid code), return unchanged
+        if not new_arms:
+            return node
+
+        # If only one arm remains and it's a wildcard, simplify to just the body
+        if len(new_arms) == 1:
+            single_arm = new_arms[0]
+            if self._is_catch_all_arm(single_arm):
+                self.stats.simplified_matches += 1
+                # Return the body directly
+                return single_arm.body
+
+        # Check if arms changed
+        if len(new_arms) != len(node.arms):
+            return nodes.MatchExprNode(
+                scrutinee=node.scrutinee,
+                arms=tuple(new_arms),
+                ensure_exhaustiveness=node.ensure_exhaustiveness,
+                source_location=node.source_location,
+            )
+
+        return node
+
+    def _simplify_arm_guard(self, arm: nodes.MatchArm) -> nodes.MatchArm:
+        """
+        Simplify arm guard if it's a constant TRUE.
+
+        A guard that's always TRUE can be removed entirely.
+        """
+        if arm.guard is None:
+            return arm
+
+        # Check if guard is constant TRUE
+        guard = arm.guard
+        if self._constant_folder:
+            guard = self._constant_folder.transform(guard)
+
+        if isinstance(guard, nodes.BoolLit) and guard.value is True:
+            self.stats.removed_true_guards += 1
+            return nodes.MatchArm(
+                pattern=arm.pattern,
+                guard=None,  # Remove the always-true guard
+                body=arm.body,
+                source_location=arm.source_location,
+            )
+
+        # Guard is constant FALSE - this arm is dead
+        # (but ReachabilityChecker should have caught this)
+
+        return arm
+
+    def _is_catch_all_arm(self, arm: nodes.MatchArm) -> bool:
+        """
+        Check if an arm is a catch-all (wildcard with no guard).
+
+        A catch-all arm matches any value unconditionally.
+        """
+        if arm.guard is not None:
+            return False
+
+        pattern = arm.pattern
+        return isinstance(pattern, (nodes.WildcardPattern, nodes.BindingPattern))
+
+
+def eliminate_dead_code(
+    ast: nodes.ASTNode,
+    type_info: Optional[TypeInferenceResult] = None,
+    **kwargs,
+) -> nodes.ASTNode:
+    """
+    Convenience function to eliminate dead code from an AST.
+
+    Args:
+        ast: The AST to transform.
+        type_info: Optional type inference result for better analysis.
+        **kwargs: Options passed to DeadCodeEliminator constructor.
+
+    Returns:
+        A new AST with dead code eliminated.
+
+    Example:
+        from yuho.ast.dead_code import eliminate_dead_code
+        optimized = eliminate_dead_code(parsed_ast)
+    """
+    eliminator = DeadCodeEliminator(type_info, **kwargs)
+    return eliminator.transform(ast)