m@(CompiledMethod traits) sourceTreeOf: index
"Find the source tree corresponding to a bytecode's index."
[
  m debugMap
    do: [| :start :end :sourceTree |
           (index >= start and: [index <= end])
             ifTrue: [^ sourceTree]]
    inGroupsOf: 3.
  Nil
].

m@(CompiledMethod traits) newEmpty
"Answer a new CompiledMethod with a fresh compilation state."
[| newM |
  newM: m clone.
  newM literals: m literals newEmpty.
  newM selectors: m selectors newEmpty.
  newM code: m code newEmpty.
  newM debugMap: m debugMap newEmpty. 
  newM
].

m@(CompiledMethod traits) isExtended: code
"Answer whether the number corresponds to an extended opcode."
[(code bitAnd: 16r0F) = 16r0F].

m@(CompiledMethod traits) isJump: code
"Answer whether the number corresponds to a jump opcode."
[{16r0F. 16r1F. 16r2F} includes: code].

m@(CompiledMethod traits) indexOfNextOpcodeAfter: index
"Answer the index in the bytecode array of the next opcode occurring after
the one at the given index.
ASSUME: the bytecode at the index should be interpreted as an opcode."
[| code |
  code: (m code at: index).
  (m isExtended: code)
    ifTrue:
      [(code between: 16r0F and: 16r2F) "Whether its a jump."
	ifTrue: [index + 3] "16-bit offset for jumping."
	ifFalse: [index + 1]]
    ifFalse:
      [(code bitAnd: 16rF0) = 16rF0 "Extended immediate value."
	ifTrue: [index + 1]
	ifFalse: "Loop through the 16rFF's to get past the immediate value."
	  [| result |
	   result: index + 1.
	   [(m code at: result) >= 16r80]
	     whileTrue: [result: result + 1].
	   result + 1]]
].

m@(CompiledMethod traits) opcodeIndicesDo: block
"Iterate through the CompiledMethod's bytecodes, applying the block to each
index where an opcode should be."
[| index |
  index: 0.
  [index < m code size]
    whileTrue:
      [block applyWith: index.
       index: (m indexOfNextOpcodeAfter: index)]
].

m@(CompiledMethod traits) immediateCodeAt: index
"Reads the bytecode at the index and determines first that it needs an
immediate value, and then reads in the value according to the scheme of
encoding."
[| code result |
  code: (m code at: index).
  (m isExtended: code)
    ifTrue: [error: 'The bytecode ' ; (code as: String) ; ' cannot have an immediate value.'].
  result: (code bitShift: -4).
  curIndex: index.
  result = 16rF
    ifTrue: [^ result].
  result: 0.
  [curIndex: curIndex + 1.
   result: (result bitShift: 7) + ((m code at: curIndex) bitAnd: 16r7F).
   (m code at: curIndex) >= 16r80] whileTrue.
  result
].

VM addPrototype: #ByteCompiler derivedFrom: {Cloneable}.
"A ByteCompiler of VM Bytecode CompiledMethods."
VM ByteCompiler addSlot: #contexts valued: Stack newEmpty.
"The current lexical context stack."
VM ByteCompiler traits addDelegate: #byteCodes valued: VM ByteCode.
"So the generator knows what bytecodes it supports."

g@(VM ByteCompiler traits) newEmpty
[| newG |
  newG: g clone.
  newG contexts: g contexts newEmpty.
  newG
].

VM ByteCompiler addPrototype: #Context derivedFrom: {Cloneable}.
VM ByteCompiler Context addSlot: #method valued: CompiledMethod newEmpty.
"The method the context targets."
VM ByteCompiler Context addSlot: #codeWriter.
"The stream for writing bytecodes to the method."
VM ByteCompiler Context addSlot: #selectors valued: ExtensibleArray newEmpty.
"The gathering collection for the method's literal array; must be indexable
before committing."
VM ByteCompiler Context addSlot: #literals valued: ExtensibleArray newEmpty.
"The gathering collection for the method's selector array; must be indexable
before committing."
VM ByteCompiler Context addSlot: #debugMap valued: ExtensibleArray newEmpty.
"The gathering collection for the method's debug map; must be indexable at
all times."
VM ByteCompiler Context addSlot: #labels valued: ExtensibleArray newEmpty.
"Maps from label serials to bytecode indices."
VM ByteCompiler Context addSlot: #relocations valued: Dictionary newEmpty.
"Maps from bytecode indices of jump opcodes to label serials."

c@(VM ByteCompiler Context traits) newOn: method
"Answer a new ByteCompiler Context that targets the given method. The method is
taken as-is and not cloned."
[| newC |
  newC: c clone.
  newC method: method.
  newC codeWriter: method code writer.
  newC selectors: (method selectors as: c selectors).
  newC literals: (method literals as: c literals).
  newC debugMap: (method debugMap as: c debugMap).
  newC labels: c labels newEmpty.
  newC relocations: c relocations newEmpty.
  newC
].

c@(VM ByteCompiler Context traits) copy
[| newC |
  newC: resend.
  newC selectors: c selectors copy.
  newC literals: c literals copy.
  newC labels: c labels copy.
  newC relocations: c relocations copy.
  newC
].

c@(VM ByteCompiler Context traits) resolveLabels
"This takes the labels Dictionary and uses it to replace the temporary ID's
placed in the jump fields with actual offsets."
[| code |
  code: c method code.
  c relocations keysAndValuesDo:
    [| :index :label offset |
     offset: (c labels at: label) - index - 2.
     code at: index put: offset.
     code at: index + 1 put: (offset bitShift: -8)]
].

c@(VM ByteCompiler Context traits) flush
"Act like a WriteStream, making the target method assume the final resulting
arrays necessary for it to operate."
[
  c method code: c codeWriter contents.
  c method literals: (c literals as: Array).
  c method selectors: (c selectors as: Array).
  c method debugMap: (c debugMap as: Array).
  c resolveLabels.
  c
].

gen@(VM ByteCompiler traits) method
"The current method."
[gen contexts top method].

gen@(VM ByteCompiler traits) codeWriter
"The writestream onto the current method's code ByteArray."
[gen contexts top codeWriter].

gen@(VM ByteCompiler traits) mapTo: sourceTree
[| index entry debugMap |
  index: gen codeWriter position.
  debugMap: gen contexts top debugMap.
  (debugMap isEmpty not and: [debugMap last == sourceTree])
    ifTrue: 
      [debugMap at: debugMap size - 1 put: index]
    ifFalse:
      [debugMap addAllLast: { index. index. sourceTree }]
].

gen@(VM ByteCompiler traits) emitByte: value from: node
[
  gen mapTo: node.
  gen codeWriter nextPut: value
].

gen@(VM ByteCompiler traits) emitShort: value from: node
[
  gen emitByte: (value bitAnd: 16rFF) from: node.
  gen emitByte: (value bitShift: -8) from: node
].

gen@(VM ByteCompiler traits) emitImmediate: value from: node
[| bytes |
  bytes: 0.
  [(1 bitShift: 7 * bytes) <= value]
    whileTrue: [bytes: bytes + 1].
  [bytes > 1]
    whileTrue:
      [bytes: bytes - 1.
        gen emitByte: ((value bitShift: -7 * bytes) bitAnd: 16r7F) + 16r80 from: node].
  gen emitByte: (value bitAnd: 16r7F) from: node
].

gen@(VM ByteCompiler traits) newLabel
[| label |
  label: gen contexts top labels size.
  gen contexts top labels addLast: Nil.
  label
].

gen@(VM ByteCompiler traits) emitLabel
[| label |
  label: gen newLabel.
  gen emitLabel: label.
  label
].

gen@(VM ByteCompiler traits) emitLabel: label
[
  gen contexts top labels at: label put: gen codeWriter position
].

gen@(VM ByteCompiler traits) emitRelocationAgainst: label
[
  gen contexts top relocations at: gen codeWriter position put: label
].

gen@(VM ByteCompiler traits) emitBranchTo: label from: msg
[
  gen emitRelocationAgainst: label.
  gen emitShort: 0 from: msg
].

gen@(VM ByteCompiler traits) emitLiteral: obj
"Ensure that the literal object is included in the literals array. If it is
not present already, it is appended to the end, thus ensuring that no other
indices require shifting. Answer the literal's index."
[
  gen contexts top literals include: obj.
  gen contexts top literals indexOf: obj
].

gen@(VM ByteCompiler traits) emitSelector: selector
"Ensure that the literal object is included in the selectors array. If it is
not present already, it is appended to the end, thus ensuring that no other
indices require shifting. Answer the selector's index."
[
  gen contexts top selectors include: selector.
  gen contexts top selectors indexOf: selector
].

gen@(VM ByteCompiler traits) emitInstruction: code from: node
"Emitting an instruction without an immediate value required just puts the
byte onto the end."
[
  gen emitByte: code from: node
].

gen@(VM ByteCompiler traits) emitInstruction: code withImmediate: value from: node
"For instructions taking an immediate-value argument, handle the encoding
scheme for the caller."
[
  value < 16rF
    ifTrue: [gen emitByte: code + (value bitShift: 4) from: node]
    ifFalse:
      [
        gen emitByte: (code bitOr: 16rF0) from: node.
        gen emitImmediate: value from: node
      ]
].

_@(VM ByteCompiler traits) generate: _@(Syntax Node traits)
"Do nothing in the default case, for comments and annotations and such."
[].

gen@(VM ByteCompiler traits) generate: ann@(Syntax Annotation traits)
"Generate the annotation's value."
[
  gen generate: ann value
].

gen@(VM ByteCompiler traits) generate: block@(Syntax Block traits)
"Encountering a new block, build a new CompiledMethod object and push it and
a new bytecode array writer onto the generator, then go through the underlying
code and generate that. When done, pop both, set up the block as a literal
and push it onto the stack."
[| newBlock codeWriter |
  newBlock: CompiledMethod newEmpty.
  newBlock inputVariables: block inputVariables size.
  newBlock localVariables: block localVariables size.
  newBlock restVariable: block restVariable isNotNil.
  newBlock optionalKeywords: block optionalKeywords.
  gen contexts isEmpty
    ifTrue: [newBlock environment: block parentScope namespace]
    ifFalse: [newBlock environment: gen method environment].
  newBlock sourceTree: block.
  gen contexts push: (gen Context newOn: newBlock).
  block statements size > 1
    ifTrue: [block statements allButLastDo:
          [| :node | gen generate: node].
         gen emitInstruction: gen popStack withImmediate: block statements size - 1 from: block.
         gen generate: block statements last]
    ifFalse: [gen generate: block statements first].
  gen contexts pop flush.
  "Forces the newBlock to record all the remaining stream input correctly."
  gen contexts isEmpty
    ifFalse:
      [gen emitInstruction: gen newBlock withImmediate: (gen emitLiteral: newBlock) from: block].
  newBlock
].

gen@(VM ByteCompiler traits) generate: def@(Syntax MethodDefinition traits)
"Translate method definitions to equivalent asMethod:on: invocations."
[
  gen contexts isEmpty
    ifTrue: [^ resend].
  "If the generate: happened at the top level, generate it as a normal block instead."
  gen emitInstruction: gen loadSelector withImmediate: (gen emitSelector: #asMethod:on:) from: def.
  resend.
  gen emitInstruction: gen loadSelector withImmediate: (gen emitSelector: def selector) from: def.
  def roles do: [| :role | gen generate: role].
  gen emitInstruction: gen newArray withImmediate: def roles size from: def.
  gen emitInstruction: gen sendMessage withImmediate: 3 from: def
].

gen@(VM ByteCompiler traits) generate: r@(Syntax Resend traits)
[
  0 below: gen method inputVariables
    do: [| :index |
      gen emitInstruction: gen loadVariable withImmediate: index from: r].
  gen emitInstruction: gen resendMessage from: r
].

gen@(VM ByteCompiler traits) generate: r@(Syntax Return traits)
[
  overrideThis
].

gen@(VM ByteCompiler traits) generate: r@(Syntax ReturnClose traits)
[| lexOffset scope |
  lexOffset: 0.
  scope: gen method sourceTree.
  [(scope isSameAs: Syntax MethodDefinition)
    or: [scope parentScope isSameAs: Syntax Namespace]]
    whileFalse:
      [(gen contexts fromTop: lexOffset) method sourceTree == scope
         ifTrue: [lexOffset: lexOffset + 1].
        scope: scope parentScope].
  (scope isSameAs: Syntax MethodDefinition)
    ifFalse: [error: '^ must be used within a method definition.'].
  gen generate: r value.
  gen emitInstruction: gen returnFrom withImmediate: lexOffset from: r
].

gen@(VM ByteCompiler traits) generate: r@(Syntax ReturnFar traits)
[| lexOffset scope methodOffset |
  lexOffset: 0.
  scope: gen method sourceTree.
  [scope isSameAs: Syntax Namespace]
    whileFalse:
      [(scope isSameAs: Syntax MethodDefinition)
        ifTrue:
          [methodOffset: lexOffset].
        (gen contexts fromTop: lexOffset) method sourceTree == scope
          ifTrue: [lexOffset: lexOffset + 1].
        scope: scope parentScope].
  methodOffset
    ifNil: [error: '^^ must be used within a method definition.'].
  gen generate: r value.
  gen emitInstruction: gen returnFrom withImmediate: methodOffset from: r
].

gen@(VM ByteCompiler traits) generate: literal@(Syntax Literal traits)
"Generate the appropriate literal-push bytecode and adds the literal to the
literals array."
[| index |
  literal value
    caseOf:
      {
        Nil -> [gen emitInstruction: gen pushNil from: literal].
        True -> [gen emitInstruction: gen pushTrue from: literal].
        False -> [gen emitInstruction: gen pushFalse from: literal].
      }
    otherwise:
      [
        index: (gen emitLiteral: literal value).
        gen emitInstruction: gen loadLiteral withImmediate: index from: literal
      ]
].

gen@(VM ByteCompiler traits) generate: n@(Syntax CompoundStatement traits)
"Handle the basis case of generating the code for each statement in the node
in order, pushing the results on the stack."
[
  n statements do: [| :node | gen generate: node]
].

gen@(VM ByteCompiler traits) generate: n@(Syntax Parenthesis traits)
"With this node type, only the last result should remain on the stack, so
this naive implementation must pop all but the last."
"TODO: find some way to avoid the excessive stack work.
Maybe pop each one as it proceeds?"
[
  n statements size > 1
    ifTrue: [n statements allButLastDo:
	      [| :node | gen generate: node].
	     gen emitInstruction: gen popStack withImmediate: n statements size - 1 from: n.
	     gen generate: n statements last]
    ifFalse: [gen generate: n statements first]
].

gen@(VM ByteCompiler traits) generate: i@(Syntax ImplicitArgument traits)
"The implicit argument amounts to putting the enclosing lexical environment
onto the stack."
[
  gen emitInstruction: gen pushEnvironment from: i
].

gen@(VM ByteCompiler traits) generate: _@(Syntax Namespace traits)
[
  shouldNotImplement
].

gen@(VM ByteCompiler traits) generate: load@(Syntax LoadVariable traits)
[| scope block lexOffset varIndex |
  scope: load variable scope.
  varIndex: (scope localVariables indexOf: load variable).
  block: gen method sourceTree.
  lexOffset: 0.
  [block == scope]
    whileFalse:
      [
        (gen contexts fromTop: lexOffset) method sourceTree == block
          ifTrue: [lexOffset: lexOffset + 1].
        block: block parentScope
      ].
  lexOffset = 0
    ifTrue:
      [gen emitInstruction: gen loadVariable withImmediate: varIndex from: load]
    ifFalse:
      [gen emitInstruction: gen loadFreeVariable withImmediate: lexOffset from: load.
       gen emitImmediate: varIndex from: load]
].

gen@(VM ByteCompiler traits) generate: store@(Syntax StoreVariable traits)
[| scope block lexOffset varIndex |
  gen generate: store value.
  scope: store variable scope.
  varIndex: (scope localVariables indexOf: store variable).
  block: gen method sourceTree.
  lexOffset: 0.
  [block == scope]
    whileFalse:
      [
        (gen contexts fromTop: lexOffset) method sourceTree == block
          ifTrue: [lexOffset: lexOffset + 1].
        block: block parentScope
      ].
  lexOffset = 0
    ifTrue:
      [gen emitInstruction: gen storeVariable withImmediate: varIndex from: store]
    ifFalse:
      [gen emitInstruction: gen storeFreeVariable withImmediate: lexOffset from: store.
       gen emitImmediate: varIndex from: store]
].

gen@(VM ByteCompiler traits) generate: array@(Syntax Array traits)
"Generate the code to push the element expression results on the stack,
then the appropriate literal-array constructor bytecode."
[
  resend.
  gen emitInstruction: gen newArray withImmediate: array size from: array
].

gen@(VM ByteCompiler traits) generate: selector@(Symbol traits) on: args from: msg@(Syntax Message traits)
"Generate the code to push the argument expression results on the stack, then
the push for the selector, and then the appropriate message send bytecode."
[
  gen emitInstruction: gen loadSelector withImmediate: (gen emitSelector: selector) from: msg.
  args do: [| :arg | gen generate: arg].
  gen emitInstruction: gen sendMessage withImmediate: args size from: msg
].

gen@(VM ByteCompiler traits) generate: msg@(Syntax Message traits)
[
  gen generate: msg selector on: msg arguments from: msg
].

gen@(VM ByteCompiler traits) generate: selector@(Symbol traits) on: args from: opts@(Syntax OptionalKeywords traits)
"Generate the code to push the argument expression results on the stack, then
the push for the selector, and then the appropriate message send bytecode."
[
  gen emitInstruction: gen loadSelector withImmediate: (gen emitSelector: selector) from: opts.
  args do: [| :arg | gen generate: arg].
  opts keywords with: opts arguments
    do: [| :key :arg |
          gen emitInstruction: gen loadSelector withImmediate: (gen emitSelector: key) from: opts.
          gen generate: arg
        ].
  gen emitInstruction: gen newArray withImmediate: opts keywords size + opts arguments size from: opts.
  gen emitInstruction: gen sendMessageWithOptionals withImmediate: args size from: opts
].

gen@(VM ByteCompiler traits) generate: opts@(Syntax OptionalKeywords traits)
[
  gen generate: opts message selector on: opts message arguments from: opts
].

gen@(VM ByteCompiler traits) generate: _@#True on: args from: msg
[
  gen emitInstruction: gen pushTrue from: msg
].

gen@(VM ByteCompiler traits) generate: _@#False on: args from: msg
[
  gen emitInstruction: gen pushFalse from: msg
].

gen@(VM ByteCompiler traits) generate: _@#Nil on: args from: msg
[
  gen emitInstruction: gen pushNil from: msg
].

gen@(VM ByteCompiler traits) generate: _@#== on: args from: msg
"Handles the special case of sending ==, which should have only one
implementation."
[
  args do: [| :arg | gen generate: arg].
  gen emitInstruction: gen isIdenticalTo from: msg
].

gen@(VM ByteCompiler traits) generate: _@#do on: args from: msg
"If the block is a literal with no variables, then inline it. Otherwise, fall back to evaluating
it via 'do'."
[| block |
  block: args first.
  ((block isSameAs: Syntax Block)
    and: [block localVariables size = 0])
    ifFalse: [^ resend].
  block statements size > 1
    ifTrue: [block statements allButLastDo:
          [| :node | gen generate: node].
         gen emitInstruction: gen popStack withImmediate: block statements size - 1 from: block.
         gen generate: block statements last]
    ifFalse: [gen generate: block statements first]
].

gen@(VM ByteCompiler traits) generate: _@#loop on: args from: msg
"Repeatedly evaluates a block via 'do'."
[| label |
  label: gen emitLabel.
  gen generate: #do on: args from: msg.
  gen emitInstruction: gen popStack withImmediate: 1 from: msg.
  gen emitInstruction: gen jumpTo from: msg.
  gen emitBranchTo: label from: msg.
  gen emitInstruction: gen pushNil from: msg
].

gen@(VM ByteCompiler traits) generate: _@#ifTrue:ifFalse: on: args from: msg
"Branches to one of two blocks and evaluates it via 'do'."
[| falseLabel endLabel |
  falseLabel: gen newLabel.
  endLabel: gen newLabel.
  gen generate: args first.
  gen emitInstruction: gen branchIfFalse from: msg.
  gen emitBranchTo: falseLabel from: msg.
  gen generate: #do on: {args second} from: msg.
  gen emitInstruction: gen jumpTo from: msg.
  gen emitBranchTo: endLabel from: msg.
  gen emitLabel: falseLabel.
  gen generate: #do on: {args third} from: msg.
  gen emitLabel: endLabel
].

gen@(VM ByteCompiler traits) generate: _@#ifTrue: on: args from: msg
"Optionally evaluates a block via 'do'."
[| falseLabel endLabel |
  falseLabel: gen newLabel.
  endLabel: gen newLabel.
  gen generate: args first.
  gen emitInstruction: gen branchIfFalse from: msg.
  gen emitBranchTo: falseLabel from: msg.
  gen generate: #do on: {args second} from: msg.
  gen emitInstruction: gen jumpTo from: msg.
  gen emitBranchTo: endLabel from: msg.
  gen emitLabel: falseLabel.
  gen emitInstruction: gen pushNil from: msg.
  gen emitLabel: endLabel
].

gen@(VM ByteCompiler traits) generate: _@#ifFalse: on: args from: msg
"Optionally evaluates a block via 'do'."
[| falseLabel endLabel |
  falseLabel: gen newLabel.
  endLabel: gen newLabel.
  gen generate: args first.
  gen emitInstruction: gen branchIfTrue from: msg.
  gen emitBranchTo: falseLabel from: msg.
  gen generate: #do on: {args second} from: msg.
  gen emitInstruction: gen jumpTo from: msg.
  gen emitBranchTo: endLabel from: msg.
  gen emitLabel: falseLabel.
  gen emitInstruction: gen pushNil from: msg.
  gen emitLabel: endLabel
].

gen@(VM ByteCompiler traits) generate: _@#whileTrue: on: args from: msg
"Repeatedly evaluates a block via 'do' while it returns True."
[| label endLabel |
  endLabel: gen newLabel.
  label: gen emitLabel.
  gen generate: #do on: {args first} from: msg.
  gen emitInstruction: gen branchIfFalse from: msg.
  gen emitBranchTo: endLabel from: msg.
  gen generate: #do on: {args second} from: msg.
  gen emitInstruction: gen popStack withImmediate: 1 from: msg.
  gen emitInstruction: gen jumpTo from: msg.
  gen emitBranchTo: label from: msg.
  gen emitLabel: endLabel.
  gen emitInstruction: gen pushNil from: msg
].

gen@(VM ByteCompiler traits) generate: _@#whileFalse: on: args from: msg
"Repeatedly evaluates a block via 'do' while it returns False."
[| label endLabel |
  endLabel: gen newLabel.
  label: gen emitLabel.
  gen generate: #do on: {args first} from: msg.
  gen emitInstruction: gen branchIfTrue from: msg.
  gen emitBranchTo: endLabel from: msg.
  gen generate: #do on: {args second} from: msg.
  gen emitInstruction: gen popStack withImmediate: 1 from: msg.
  gen emitInstruction: gen jumpTo from: msg.
  gen emitBranchTo: label from: msg.
  gen emitLabel: endLabel.
  gen emitInstruction: gen pushNil from: msg
].

gen@(VM ByteCompiler traits) generate: _@#whileTrue on: args from: msg
"Repeatedly evaluates a block via 'do' while it returns True."
[| label |
  label: gen emitLabel.
  gen generate: #do on: args from: msg.
  gen emitInstruction: gen branchIfTrue from: msg.
  gen emitBranchTo: label from: msg.
  gen emitInstruction: gen pushNil from: msg
].

gen@(VM ByteCompiler traits) generate: _@#whileFalse on: args from: msg
"Repeatedly evaluates a block via 'do' while it returns False."
[| label |
  label: gen emitLabel.
  gen generate: #do on: args from: msg.
  gen emitInstruction: gen branchIfFalse from: msg.
  gen emitBranchTo: label from: msg.
  gen emitInstruction: gen pushNil from: msg
].

gen@(VM ByteCompiler traits) generate: mode@(Syntax Mode traits)
[
  error: 'Modes are not supported for bytecode compilation.'
].