Rust Constant Propagation

Constant propagation is a compiler optimization technique that replaces uses of constants with their actual values. This can simplify code, enable further optimizations, and improve performance. This challenge asks you to implement a simplified form of constant propagation for a small, abstract syntax tree (AST) representing Rust-like expressions.

Problem Description

Your task is to implement a function that takes an AST representing simple arithmetic expressions and performs constant propagation. This means identifying variables that are assigned a constant value and then replacing subsequent uses of those variables with their constant values.

Key Requirements:

1. Identify Constant Assignments: Detect when a variable is assigned a literal integer value.
2. Track Variable Values: Maintain a mapping of variables to their assigned constant values.
3. Propagate Constants: When a variable is used after being assigned a constant, replace the variable node with the constant node in the AST.
4. Handle Reassignment: If a variable is reassigned, update its tracked value and propagate the new value.
5. Handle Non-Constant Assignments: Ignore assignments to variables that are not literal integers or that depend on other variables.
6. Handle Scope (Simplified): For this challenge, assume a single, global scope. Variable assignments are visible and affect all subsequent uses.

Expected Behavior:

The function should traverse the AST. When it encounters an assignment to a variable with a literal integer, it should record this mapping. When it encounters an expression that uses a variable, if that variable has a known constant value, it should replace the variable node in the AST with the corresponding constant node. Operations involving variables that have been replaced by constants should be evaluated if possible, further simplifying the AST.

Edge Cases to Consider:

• Variables used before they are assigned a constant value.
• Assignments to variables that are not literal integers (e.g., x = y + 1).
• Expressions involving operations where at least one operand is a propagated constant.

Examples

Example 1:

Input AST:

Block {
    statements: [
        Assignment {
            variable: "a",
            value: Literal(10)
        },
        Expression {
            value: BinaryOp {
                left: Variable("a"),
                op: Add,
                right: Literal(5)
            }
        }
    ]
}

Output AST:

Block {
    statements: [
        Assignment {
            variable: "a",
            value: Literal(10)
        },
        Expression {
            value: Literal(15) // 10 + 5 is evaluated
        }
    ]
}

Explanation: The variable a is assigned the literal value 10. In the subsequent expression, a is replaced by 10, and the addition 10 + 5 is evaluated to 15.

Example 2:

Input AST:

Block {
    statements: [
        Assignment {
            variable: "x",
            value: Literal(20)
        },
        Assignment {
            variable: "y",
            value: Variable("x")
        },
        Expression {
            value: Variable("y")
        }
    ]
}

Output AST:

Block {
    statements: [
        Assignment {
            variable: "x",
            value: Literal(20)
        },
        Assignment {
            variable: "y",
            value: Literal(20) // y is assigned x, and x is 20
        },
        Expression {
            value: Literal(20) // y is replaced by its propagated value
        }
    ]
}

Explanation: x is assigned 20. y is assigned x. Since x is 20, y effectively becomes 20. The final expression uses y, which is then replaced by 20.

Example 3 (Reassignment):

Input AST:

Block {
    statements: [
        Assignment {
            variable: "count",
            value: Literal(1)
        },
        Expression {
            value: Variable("count")
        },
        Assignment {
            variable: "count",
            value: Literal(5)
        },
        Expression {
            value: Variable("count")
        }
    ]
}

Output AST:

Block {
    statements: [
        Assignment {
            variable: "count",
            value: Literal(1)
        },
        Expression {
            value: Literal(1)
        },
        Assignment {
            variable: "count",
            value: Literal(5)
        },
        Expression {
            value: Literal(5)
        }
    ]
}

Explanation: count is first assigned 1. The first expression uses count, which is propagated as 1. Then, count is reassigned to 5. The second expression uses count, which is now propagated as 5.

Constraints

• The AST will contain only Block, Assignment, Expression, Variable, Literal, and BinaryOp nodes.
• Literal nodes will only contain i32 integer values.
• BinaryOp will only involve Add, Sub, Mul, Div.
• Variable names will be alphanumeric strings.
• The AST structure will be valid (e.g., no dangling variables in assignments).
• Performance is not a critical concern for this challenge, but an efficient traversal is desirable.

Notes

You will need to define the AST data structures in Rust. Consider how to represent the AST immutably or mutably during the propagation process. Think about the order of operations: should you propagate first and then evaluate, or can these be interleaved? A common approach is to use a recursive function that traverses the AST and returns a potentially modified AST node, along with the current state of variable bindings. You may find it useful to use a HashMap to store variable bindings.