link.hy

(import dataclasses [dataclass])
(import json)
(import os)
(import struct)
(import sys)

(import hy)
(import hy.models [Expression Integer Symbol])
(require hyrule.control [defmain unless])
(require hyrule.misc [of pun])
(import hyrule.hypprint [pprint])

(import models [CompiledFunction LinkedFunction Program Unit])


;; Information about a function, necessary and sufficient to link against it
;; Normally comes from units, but can also be injected by REPL
(defclass [dataclass] ProgramFunction []
  #^ int id
  #^ str name
  #^ int argc
  #^ int retc
  )

;; Additional information not included in program or program fragment
(defclass [dataclass] LinkInfo []
  #^ int bc-end
  #^ (of dict str ProgramFunction) function-table
  #^ (of dict str int) global-table
  )

(defn link-program [units
                    output
                    builtin-functions
                    [repl-initial-state None]
                    [allow-no-main False]]
  (setv bc-end 0)

  (setv #^ (of dict str ProgramFunction)
        function-table {})
  (setv #^ (of dict str int)
        global-table {})

  (setv functions-to-compile [])

  (setv program (Program :bytecode []
                         :functions []
                         :globals []))

  (when (is-not repl-initial-state None)
    (setv bc-end          repl-initial-state.bc-end
          function-table  {#** repl-initial-state.function-table}
          global-table    {#** repl-initial-state.global-table}))

  ;; link:
  ;; - collect functions + globals

  (for [unit units]
    (for [f unit.functions]
      (when (in f.name function-table)
        (raise (Exception f"Multiple definitions of function '{f.name}'")))

      (setv function-index (len function-table))
      (setv (get function-table f.name)
            (ProgramFunction :id function-index
                             :name f.name
                             :argc f.argc
                             :retc f.retc))
      (functions-to-compile.append f)
      )
    (for [#(g value) (unit.globals.items)]
      (if (is value None)
          ;; If a global has no value, it has been declared akin to 'extern' in C.
          ;; The compiler doesn't currently permit this, but variables in the REPL use this.
          (get global-table g)
          (do
            (when (in g global-table)
              (raise (Exception f"Multiple definitions of global variable '{g}'")))

            ;; Note: global-table may contain additional entries coming from repl-initial-state,
            ;;       which will not be found in program.globals (and must not be added there)
            (setv global-index (len global-table))
            (setv (get global-table g) global-index)
            (program.globals.append value)
            ))
      )
    )

  ;; resolve function calls
  ;; resolve global var IDs

  (defn error [message]
    (raise (Exception f"error: {message}")))

  (for [f functions-to-compile]
    (defn resolve [insn]
      (cond
        ;; call
        (= (get insn 0) 'call) (do
          (setv [_ name argc retc] insn)

          (defn check-retc [produces]
            (unless (= produces retc)
              (error f"{retc} results were expected, but function '{name}' produces {produces}")))

          (try
            (setv f (get function-table name))
            (except [KeyError]
              (raise (Exception f"unresolved function {name}") :from None)))

          (unless (= argc f.argc)
            (error f"Function '{name}' expects {f.argc} arguments, but {argc} were passed"))
          (check-retc f.retc)
          ['call f.id])
        ;; getglobal/setglobal
        (in (get insn 0) #{'getglobal 'setglobal}) (do
          (setv [opcode name] insn)
          [opcode (get global-table name)])

        True insn
        ))

    (setv f.body (lfor insn f.body (resolve insn)))
    )

  ;; expand jump offsets to bytes

  (defn instruction-length [insn]
    ;; 1 byte per opcode and each operand
    ;; except for branches & pushconst where the operand is 2 bytes
    (if (in (get insn 0) #{'jmp 'jz 'pushconst})
      3
      (len insn)))

  (for [f functions-to-compile]
    (defn resolve [i insn]
      (cond
        ;; getglobal/setglobal
        (in (get insn 0) #{'jmp 'jz}) (do
          (setv [opcode dist] insn)

          ;; distance is sum of lengths of instructions, starting at the next one
          (setv start (+ i 1))
          (setv end (+ start dist))

          (defn block-length [block]
            (sum (gfor insn block (instruction-length insn))))

          (if (>= end start)
            (setv dist-bytes (block-length (cut f.body start end)))
            (setv dist-bytes (- (block-length (cut f.body end start))))
            )

          [opcode dist-bytes])

        True insn
        ))

    (setv f.body (lfor [i insn] (enumerate f.body) (resolve i insn)))
    )


  ;; - lay out bytecode

  (for [f functions-to-compile]
    (setv func-body-len
          (sum (gfor insn f.body (instruction-length insn))))

    (setv f* (LinkedFunction :name f.name
                             :argc f.argc
                             :num-locals f.num-locals
                             :bytecode-offset bc-end))
    (program.functions.append f*)

    (setv program.bytecode (+ program.bytecode f.body)
          bc-end (+ bc-end func-body-len))
    )


  ;(pprint all-functions)
  (pprint global-table)
  ;; (pprint function-table)
  (pprint program)

  (setv OPCODE-NUMBERS {
    'pushconst 0
    'zero 1
    'drop 2
    'getglobal 3
    'setglobal 4
    'getlocal 5
    'setlocal 6
    'call 10
    'ret 13
    'jmp 20
    'jz 21
    })

  (when (is-not output None)
    (if allow-no-main
      ;; In REPL mode, definitions generate program fragments with no main function
      (setv main-func-idx (try (. function-table ["main"] id) (except [KeyError] 255)))
      (setv main-func-idx (. function-table ["main"] id)))

    (with [f (open (+ output ".tmp") "wb")]
      ;; helper function for writing binary data
      (defn emit [format #* args]
        (f.write (struct.pack format #* args)))

      ;; write header
      (emit "<HBBBxxx"
            bc-end
            (len program.functions)
            (len program.globals)
            main-func-idx)
      ;; functions
      (for [func program.functions]
        (emit "<BBH" func.argc func.num-locals func.bytecode-offset))
      ;; globals
      (for [value program.globals]
        (emit "<h" value))
      ;; bytecode
      (for [[opcode #* operands] program.bytecode]
        ;(f.write (bytes [(get OPCODE-NUMBERS opcode) #* operands])))
        (cond
          (in (str opcode) builtin-functions) (do
            (assert (= (len operands) 0))
            (emit "B" (get (get builtin-functions (str opcode)) "opcode")))
          (in opcode #{'jmp 'jz 'pushconst}) (do
            ;; branch instructions & pushconst have a 16-bit operand
            (emit "b" (get OPCODE-NUMBERS opcode))
            (emit "h" #* operands))
          True (do
            (for [b [(get OPCODE-NUMBERS opcode) #* operands]]
              (emit "B" b)))))
    )

    (os.rename (+ output ".tmp") output))
  (pun (LinkInfo :!bc-end :!function-table :!global-table)))

(defmain []
  (import argparse [ArgumentParser])

  (setv parser (ArgumentParser))
  (parser.add-argument "inputs" :nargs "+")
  (parser.add-argument "-o" :dest "output" :required True)
  (setv args (parser.parse-args))

  (with [f (open "builtins.json")]
    (setv builtin-functions (json.load f))
    )

  (setv units [])

  (for [path args.inputs]
    (with [f (open path)]
      (setv [form] (hy.read-many f))
      (setv unit (Unit.from-form form))
      ;; (pprint unit)
      (units.append unit)
      )
    )

  (link-program units
                :output args.output
                :builtin-functions builtin-functions
                )
  )