L3.uniqify

  1from collections.abc import Callable, Mapping
  2from functools import partial
  3
  4from util.sequential_name_generator import SequentialNameGenerator
  5
  6from .syntax import (
  7    Abstract,
  8    Allocate,
  9    Apply,
 10    Begin,
 11    Branch,
 12    Identifier,
 13    Immediate,
 14    Let,
 15    LetRec,
 16    Load,
 17    Primitive,
 18    Program,
 19    Reference,
 20    Store,
 21    Term,
 22)
 23
 24type Context = Mapping[Identifier, Identifier]
 25
 26
 27# This pass is responsible for renaming all identifiers in a program to be unique. This is necessary for the later stages of the compiler, which rely on the fact that all identifiers are unique.
 28def uniqify_term(
 29    term: Term,
 30    context: Context,
 31    fresh: Callable[[str], str],
 32) -> Term:
 33    _term = partial(uniqify_term, context=context, fresh=fresh)
 34
 35    match term:
 36        case Let(bindings=bindings, body=body):
 37            new_bindings: list[tuple[Identifier, Term]] = []
 38            new_context = context
 39            for name, value in bindings:
 40                new_name = fresh(name)
 41                new_bindings.append((new_name, _term(value)))
 42                new_context = {**new_context, name: new_name}
 43
 44            return Let(
 45                bindings=new_bindings,
 46                body=uniqify_term(body, new_context, fresh),
 47            )
 48
 49        case LetRec(bindings=bindings, body=body):
 50            new_bindings: list[tuple[Identifier, Term]] = []
 51            new_context = context
 52            for name, value in bindings:
 53                new_name = fresh(name)
 54                new_bindings.append((new_name, value))
 55                new_context = {**new_context, name: new_name}
 56
 57            _new_term = partial(uniqify_term, context=new_context, fresh=fresh)
 58
 59            return LetRec(
 60                bindings=[(new_name, _new_term(value)) for new_name, value in new_bindings],
 61                body=_new_term(body),
 62            )
 63
 64        case Reference(name=name):
 65            if name in context:
 66                return Reference(name=context[name])
 67            return term
 68
 69        case Abstract(parameters=parameters, body=body):
 70            new_parameters = [fresh(parameter) for parameter in parameters]
 71            new_context = {
 72                **context,
 73                **{parameter: new_parameter for parameter, new_parameter in zip(parameters, new_parameters)},
 74            }
 75            return Abstract(
 76                parameters=new_parameters,
 77                body=uniqify_term(body, new_context, fresh),
 78            )
 79
 80        case Apply(target=target, arguments=arguments):
 81            return Apply(
 82                target=_term(target),
 83                arguments=[_term(argument) for argument in arguments],
 84            )
 85
 86        case Immediate():
 87            return term
 88
 89        case Primitive(operator=operator, left=left, right=right):
 90            return Primitive(
 91                operator=operator,
 92                left=_term(left),
 93                right=_term(right),
 94            )
 95
 96        case Branch(operator=operator, left=left, right=right, consequent=consequent, otherwise=otherwise):
 97            return Branch(
 98                operator=operator,
 99                left=_term(left),
100                right=_term(right),
101                consequent=_term(consequent),
102                otherwise=_term(otherwise),
103            )
104
105        case Allocate():
106            return term
107
108        case Load(base=base, index=index):
109            return Load(
110                base=_term(base),
111                index=index,
112            )
113
114        case Store(base=base, index=index, value=value):
115            return Store(
116                base=_term(base),
117                index=index,
118                value=_term(value),
119            )
120
121        case Begin(effects=effects, value=value):  # pragma: no branch
122            return Begin(
123                effects=[_term(effect) for effect in effects],
124                value=_term(value),
125            )
126
127
128def uniqify_program(
129    program: Program,
130) -> tuple[Callable[[str], str], Program]:
131    fresh = SequentialNameGenerator()
132
133    _term = partial(uniqify_term, fresh=fresh)
134
135    match program:
136        case Program(parameters=parameters, body=body):  # pragma: no branch
137            local = {parameter: fresh(parameter) for parameter in parameters}
138            return (
139                fresh,
140                Program(
141                    parameters=[local[parameter] for parameter in parameters],
142                    body=_term(body, local),
143                ),
144            )