optimizer.py

"""
Code Optimizer
Performs optimization on Three-Address Code (TAC) generated by the ICG.
Applies multiple optimization techniques to improve efficiency and remove redundancy.

Optimizations implemented:
- Constant Folding
- Constant Propagation
- Copy Propagation
- Common Subexpression Elimination
- Dead Code Elimination
"""

from ir import Instr 

#Constant folding + propagation (combined pass)
def is_const(x):
    if x is None:
        return False
    try:
        float(x)
    except ValueError:
        return False
    return True

def parse_num(x):
    if "." in x:
        return float(x)
    return int(x)

def format_num(value):
    if isinstance(value, float) and value.is_integer():
        return str(int(value))
    return str(value)

def eval_binary(op, left, right):
    a = parse_num(left)
    b = parse_num(right)
    if op == "+":
        return format_num(a + b)
    if op == "-":
        return format_num(a - b)
    if op == "*":
        return format_num(a * b)
    if op == "/":
        if b == 0:
            return None
        return format_num(a / b)
    if op == "%":
        if b == 0:
            return None
        return format_num(a % b)
    return None

def constant_fold_and_propagate(instrs):
    const_map = {}
    for instr in instrs:
        if instr.label or instr.op in ["if", "goto", "print"]:
            if instr.op == "if":
                if instr.arg1 in const_map:
                    instr.arg1 = const_map[instr.arg1]
                if instr.arg2 in const_map:
                    instr.arg2 = const_map[instr.arg2]
            elif instr.op == "print":
                if instr.arg1 in const_map:
                    instr.arg1 = const_map[instr.arg1]
            if instr.op in ["if", "goto"] or instr.label:
                const_map = {}
            continue

        if instr.arg1 in const_map:
            instr.arg1 = const_map[instr.arg1]
        if instr.arg2 in const_map:
            instr.arg2 = const_map[instr.arg2]

        if instr.op in ["+", "-", "*", "/", "%"] and is_const(instr.arg1) and is_const(instr.arg2):
            folded = eval_binary(instr.op, instr.arg1, instr.arg2)
            if folded is not None:
                instr.op = "="
                instr.arg1 = folded
                instr.arg2 = None

        if instr.op == "=" and instr.result and is_const(instr.arg1):
            const_map[instr.result] = instr.arg1
        elif instr.result:
            const_map.pop(instr.result, None)
    return instrs

def is_barrier(instr):
    return instr.label is not None or instr.op in ["if", "goto"]

def optimize_straight_line_block(block_instrs):
    changed = True
    optimized = block_instrs
    while changed:
        before = [to_string(i) for i in optimized]
        optimized = constant_fold_and_propagate(optimized)
        optimized = copy_propagation(optimized)
        optimized = cse(optimized)
        after = [to_string(i) for i in optimized]
        changed = before != after
    return optimized

def optimize_by_blocks(instrs):
    result = []
    block = []
    for instr in instrs:
        if is_barrier(instr):
            if block:
                result.extend(optimize_straight_line_block(block))
                block = []
            result.append(instr)
            continue
        block.append(instr)
    if block:
        result.extend(optimize_straight_line_block(block))
    return result
    
#Copy propagation
def copy_propagation(instrs):
    copy_map = {}
    for instr in instrs:
        # Only substitute temporaries, never user variables
        if instr.arg1 in copy_map and is_temp(instr.arg1):
            instr.arg1 = copy_map[instr.arg1]
        if instr.arg2 in copy_map and is_temp(instr.arg2):
            instr.arg2 = copy_map[instr.arg2]

        if instr.op == "=" and instr.result and instr.arg1:
            # Only track temp = something or var = temp
            if is_temp(instr.result) and instr.arg1.isidentifier():
                copy_map[instr.result] = instr.arg1
            else:
                copy_map.pop(instr.result, None)  # user var reassigned, invalidate

    return instrs


#Common subexpression elimination
def normalize(op, a, b):
    if op in ["+", "*"]:
        return tuple(sorted([a, b]))
    return (a, b)

def cse(instrs):
    expr_map = {}

    for instr in instrs:
        if instr.op in ["+", "-", "*", "/", "%"]:
            key = (instr.op, *normalize(instr.op, instr.arg1, instr.arg2))

            if key in expr_map:
                instr.op = "="
                instr.arg1 = expr_map[key]
                instr.arg2 = None
            else:
                expr_map[key] = instr.result

    return instrs

def is_temp(var):
    """Temporaries start with 't' followed by digits: t0, t1, ..."""
    return var is not None and len(var) > 1 and var[0] == 't' and var[1:].isdigit()

def is_variable_name(token):
    return token is not None and token.isidentifier()

def dead_code(instrs):
    used = set()
    new_instrs = []
    for instr in reversed(instrs):
        keep = False
        if instr.op in ["if", "goto", "print"] or instr.label:
            keep = True
        elif instr.result and instr.result in used:
            keep = True

        if keep:
            new_instrs.append(instr)
            if instr.result and instr.op not in ["if", "goto", "print"]:
                used.discard(instr.result)
            if is_variable_name(instr.arg1):
                used.add(instr.arg1)
            if is_variable_name(instr.arg2):
                used.add(instr.arg2)

    return list(reversed(new_instrs))    

#Convert back to 3AC
def to_string(instr):
    if instr.label:
        return f"{instr.label}:"

    if instr.op == "print":
        return f"print {instr.arg1}"

    if instr.op == "if":
        return f"if {instr.arg1} {instr.relop} {instr.arg2} goto {instr.target}"

    if instr.op == "goto":
        return f"goto {instr.target}"

    if instr.op == "=":
        return f"{instr.result} = {instr.arg1}"

    if instr.arg2:
        return f"{instr.result} = {instr.arg1} {instr.op} {instr.arg2}"

    return instr.raw

#final pipeline
def optimize(lines):
    instrs = [Instr(line) for line in lines]

    instrs = optimize_by_blocks(instrs)
    instrs = dead_code(instrs)

    return [to_string(i) for i in instrs]