requires: {#Dictionary. #Stack}.
provides: {#IRGenerator}.

Optimizer addPrototype: #IRGeneratorContext derivedFrom: {Cloneable}.
"This is the part of the Memento pattern that gets saved as a roll-back state.
In this case, these slots are pushed as a separate object onto the generator's
context stack."
Optimizer IRGeneratorContext addSlot: #mode valued: Syntax Mode Slate.
Optimizer IRGeneratorContext addSlot: #method.
Optimizer IRGeneratorContext addSlot: #inputVariables.
Optimizer IRGeneratorContext addSlot: #scope.
Optimizer IRGeneratorContext addSlot: #basicBlock.
Optimizer IRGeneratorContext addSlot: #propagateConstants valued: True.

Optimizer addPrototype: #IRGenerator derivedFrom: {Optimizer IRGeneratorContext}.
Optimizer IRGenerator addSlot: #variables valued: Dictionary newEmpty.
Optimizer IRGenerator addSlot: #labels valued: Dictionary newEmpty.
Optimizer IRGenerator addSlot: #results valued: Stack newEmpty.
Optimizer IRGenerator addSlot: #controlFlow valued: Stack newEmpty.
Optimizer IRGenerator addSlot: #contexts valued: Stack newEmpty.

gen@(Optimizer IRGenerator traits) newEmpty
[| newGen |
  newGen: gen clone.
  newGen variables: gen variables newEmpty.
  newGen labels: gen labels newEmpty.
  newGen results: gen results newEmpty.
  newGen controlFlow: gen controlFlow newEmpty.
  newGen contexts: gen contexts newEmpty.
  newGen
].

gen@(Optimizer IRGenerator traits) pushContext
"Pushes the Generator's current context state as a separate Context object
onto the Generator's stack."
[| context |
  context: Optimizer IRGeneratorContext clone.
  context mode: gen mode.
  context method: gen method.
  context inputVariables: gen inputVariables.
  context scope: gen scope.
  context basicBlock: gen basicBlock.
  gen contexts push: context.
  gen
].

gen@(Optimizer IRGenerator traits) popContext
"Set the Generator's context state from the top (popped) context object from
the Generator's stack."
[| context |
  context: gen contexts pop.
  gen mode: context mode.
  gen method: context method.
  gen inputVariables: context inputVariables.
  gen scope: context scope.
  gen basicBlock: context basicBlock.
  gen
].

gen@(Optimizer IRGenerator traits) emitInstruction: instruction
"Emits an instruction, generating a basicBlock as necessary. However, this
method requires the control flow stack to be set up correctly before
being called."
[| basicBlock |
  basicBlock: gen basicBlock.
  basicBlock ifNil:
    [basicBlock: Optimizer IR BasicBlock newEmpty.
     gen emitControlFlow: basicBlock].
  basicBlock addInstruction: instruction
].

gen@(Optimizer IRGenerator traits) emitControlFlow: cf@(Optimizer IR ControlFlow traits)
[
  gen basicBlock: Nil.
  gen controlFlow top addChild: cf
].

gen@(Optimizer IRGenerator traits) emitControlFlow: bblock@(Optimizer IR BasicBlock traits)
[
  gen basicBlock: bblock.
  gen controlFlow top addChild: bblock
].

gen@(Optimizer IRGenerator traits) makeInputVariable: type
[| var |
  var: Optimizer IR InputVariable clone.
  var type: type.
  var scope: gen scope.
  gen scope ifNotNilDo: [| :scope | scope inputVariables addLast: var]
].

gen@(Optimizer IRGenerator traits) makeOutputVariable: type
[| var |
  var: Optimizer IR OutputVariable clone.
  var type: type.
  var scope: gen scope.
  gen scope ifNotNilDo: [| :scope | scope outputVariables addLast: var]
].

gen@(Optimizer IRGenerator traits) makeLocalVariable: varNode
[| var |
  var: Optimizer IR LocalVariable clone.
  var type: varNode type.
  var scope: gen scope.
  var name: varNode name.
  gen scope ifNotNilDo: [| :scope | scope localVariables add: var].
  gen variables at: varNode put: var
].

gen@(Optimizer IRGenerator traits) addLocalVariables: blockNode
[
  blockNode localVariables
    do: [| :varNode | gen makeLocalVariable: varNode]
].

gen@(Optimizer IRGenerator traits) addInputVariables: blockNode
[| set |
  blockNode inputVariables
    collect:
      [| :varNode |
        set: Optimizer IR SetVariable clone.
        set inputs: {gen makeInputVariable: varNode type}.
        set outputs: {gen makeLocalVariable: varNode}.
        gen emitInstruction: set]
].

gen@(Optimizer IRGenerator traits) makeTemporaryVariable: type
[| var |
  var: Optimizer IR LocalVariable clone.
  var type: type.
  var scope: gen scope.
  gen scope ifNotNilDo: [| :scope | scope localVariables add: var]
].

gen@(Optimizer IRGenerator traits) makeConstant: type valued: value
[| constant |
  constant: Optimizer IR Constant clone.
  ((gen mode = Syntax Mode Primitive) and: [type = Types Any])
    ifTrue: [type: value primitiveType].
  constant type: type.
  constant value: value.
  gen scope ifNotNilDo: [| :scope | scope constants add: constant]
].

gen@(Optimizer IRGenerator traits) generateArgument: argument
[
  gen propagateConstants
    ifTrue:
      [gen results push: (gen makeTemporaryVariable: argument type).
       gen generate: argument.
       gen results pop]
    ifFalse:
      [gen results push: (gen makeTemporaryVariable: argument type).
       gen pushContext.
       gen propagateConstants: True.
       gen generate: argument.
       gen popContext.
       gen results pop]
].

gen@(Optimizer IRGenerator traits) generateResult: result
[
  gen propagateConstants
    ifTrue:
      [gen pushContext.
       gen propagateConstants: False.
       gen generate: result.
       gen popContext.
       gen results top]
    ifFalse:
      [gen generate: result]
].

gen@(Optimizer IRGenerator traits) generateResult: result into: var
[
  gen propagateConstants
    ifTrue:
      [gen results push: var.
       gen pushContext.
       gen propagateConstants: False.
       gen generate: result.
       gen popContext.
       gen results pop]
    ifFalse:
      [gen results push: var.
       gen generate: result.
       gen results pop]
].

gen@(Optimizer IRGenerator traits) generate: blockNode@(Syntax Block traits)
[| closure label inst |
  closure: Optimizer IR Closure newEmpty.
  closure parent: gen scope.
  gen scope ifNotNilDo: [| :scope | scope closures add: closure].
  gen pushContext.
  gen scope: closure.
  label: Optimizer IR Continuation clone.
  closure makeChild: label.
  gen controlFlow push: label.
  gen basicBlock: Nil.
  gen addLocalVariables: blockNode.
  gen addInputVariables: blockNode.
  blockNode statements
    allButLast: 1 do: [| :statement | gen generateArgument: statement].
  gen makeOutputVariable: (blockNode type ifNotNilDo:
    [| :type | type resultType]).
  blockNode statements isEmpty
    ifFalse: [gen generateResult: blockNode statements last
		  into: closure outputVariables first].
  gen controlFlow pop. 
  gen popContext.
  gen scope
    ifNotNil:
      [
        inst: Optimizer IR InstantiateClosure clone.
        inst inputs: {closure}.
        inst outputs: {gen results top}.
        "This must be performed before the pops since it presumes stack elements."
        gen emitInstruction: inst
      ].
  closure
].

gen@(Optimizer IRGenerator traits) generate: methNode@(Syntax MethodDefinition traits)
[| meth label def |
  meth: (Optimizer IR Method newNamed: methNode selector).
  meth parent: gen scope.
  gen scope ifNotNilDo: [| :scope | scope closures add: meth].
  gen pushContext.
  gen method: meth.
  gen scope: meth.
  label: Optimizer IR Continuation clone.
  meth makeChild: label.
  gen controlFlow push: label.
  gen basicBlock: Nil.
  gen addLocalVariables: methNode.
  gen inputVariables: (gen addInputVariables: methNode).
  methNode statements
    allButLast: 1
    do: [| :statement | gen generateArgument: statement].
  gen makeOutputVariable: (methNode type ifNotNilDo:
    [| :type | type resultType]).
  methNode statements isEmpty ifFalse:
    [gen generateResult: methNode statements last into: meth outputVariables first].
  gen results pop.
  gen controlFlow pop.
  gen popContext.
  def: Optimizer IR DefineMethod clone.
  def inputs:
    "TODO: Closures/methods shouldn't be directly specified as operands?"
    {meth} ; (methNode roles
        collect: [| :role | gen generateArgument: role]).
  def outputs: {gen results top}.
  "This must be performed before the pops since it presumes stack elements."
  gen emitInstruction: def.
  meth
].

gen@(Optimizer IRGenerator traits) generate: _@(Syntax Resend traits)
[| res |
  res: Optimizer IR ResendMethod clone.
  res inputs:
    {gen makeConstant: Types Any valued: gen method selector} ;
    gen inputVariables.
  res outputs: {gen results top}.
  gen emitInstruction: res
].

gen@(Optimizer IRGenerator traits) generate: retNode@(Syntax ReturnClose traits)
[| ret |
  ret: Compiler IR Return clone.
  gen generateResult: retNode value into: gen method outputVariables first.
  ret label: gen method child.
  gen emitControlFlow: ret
].

gen@(Optimizer IRGenerator traits) generate: retNode@(Syntax ReturnFar traits)
[| ret meth |
  ret: Compiler IR Return clone.
  (gen contexts findFirst: [| :context | context method isNotNil])
    ifNotNilDo: [| :index | meth: (gen contexts at: index) method].
  gen generateResult: retNode value into: meth outputVariables first.
  ret label: meth child.
  gen emitControlFlow: ret
].

gen@(Optimizer IRGenerator traits) generate: literalNode@(Syntax Literal traits)
[| constant set |
  constant: (gen makeConstant: literalNode type valued: literalNode value).
  gen propagateConstants ifTrue: [^ constant].
  set: Optimizer IR Set clone.
  set inputs: {constant}.
  set outputs: {gen results top}.
  gen emitInstruction: set
].

gen@(Optimizer IRGenerator traits) generate: parenNode@(Syntax Parenthesis traits)
"Handles a compound statement sequence, with the expressions as inputs."
[
  parenNode statements
    allButLast: 1 do: [| :statement | gen generateArgument: statement].
  parenNode statements isEmpty
    ifFalse: [gen generate: parenNode statements last]
].

gen@(Optimizer IRGenerator traits) generate: _@(Syntax ImplicitArgument traits)
"Handles an implicit context send with a CurrentNamespace instruction."
[| namespace |
  namespace: Optimizer IR CurrentNamespace clone.
  namespace outputs: {gen results top}.
  gen emitInstruction: namespace
].

gen@(Optimizer IRGenerator traits) generate: loadNode@(Syntax LoadVariable traits)
"Generates a LoadFreeVariable."
[| load |
  load: Optimizer IR LoadFreeVariable clone.
  load inputs: {gen variables at: loadNode variable}.
  load outputs: {gen results top}.
  gen emitInstruction: load
].

gen@(Optimizer IRGenerator traits) generate: storeNode@(Syntax StoreVariable traits)
"Generates a StoreFreeVariable followed by a SetVariable."
[| store val set |
  val: (gen generateArgument: storeNode value).
  store: Optimizer IR StoreFreeVariable clone.
  store inputs: {val}.
  store outputs: {gen variables at: storeNode variable}.
  gen emitInstruction: store.
  set: Optimizer IR SetVariable clone.
  set inputs: {val}.
  set outputs: {gen results top}.
  gen emitInstruction: set
].

gen@(Optimizer IRGenerator traits) generate: arrNode@(Syntax Array traits)
"Handle a literal array, with the statements as inputs."
[| arr |
  arr: Optimizer IR CreateArray clone.
  arr inputs: (arrNode statements
        collect: [| :statement | gen generateArgument: statement]).
  arr outputs: {gen results top}.
  gen emitInstruction: arr
].

gen@(Optimizer IRGenerator traits) generate: msgNode@(Syntax Message traits)
"Force a primitive call for primitiveMode, or create a node for the send.
The selector is the first input, a constant."
[| msg |
  (gen mode = Syntax Mode Primitive)
    ifTrue: [^ (gen generatePrimitive: msgNode selector for: msgNode arguments)].
  msg: Optimizer IR InvokeMethod clone.
  msg inputs:
    ({gen makeConstant: Types Any valued: msgNode selector} ;
      (msgNode arguments
        collect: [| :argument | gen generateArgument: argument])).
  msg outputs: {gen results top}.
  gen emitInstruction: msg
].

gen@(Optimizer IRGenerator traits) generate: modeNode@(Syntax Mode traits)
"Set the mode in a new context and recurse into the modeNode's wrapper."
[
  gen pushContext.
  gen mode: modeNode mode.
  gen generate: modeNode value.
  gen popContext
].