;;; Sex semantic engine

(import
  scheme
  (chicken base)
  (chicken keyword)
  (chicken string)
  (chicken module)
  fmt
  infer
  sex-macros
  sex-modules
  types
  matchable                         ; pattern matching
  srfi-1                            ; list routines
  srfi-69                           ; hash tables
  utils
  )

(export/rename (process semen-process))
(export closure-env-declaration)

;;; for lambda extraction, docstring processing, macro expansion,
;;; injection of module headers, i.e. all things that rearrange code
;;; structurally, add or remove forms
;;;
;;; The algorithm: feed toplevel forms to appropriate handlers, then
;;; append their return to the resulting list. Each handler can return
;;; multiple forms, e.g. lambdas collected from a function may result
;;; in auxiliary structures and functions.
(define (process raw-sex-forms)
  (process-rec raw-sex-forms (list)))

(define (process-rec forms acc)
  (cond
   ((null? forms) (reverse acc))
   ((macro? (car forms))
    (process-rec
     (macroexpand (car forms) (cdr forms))
     acc))
   (else
    ;; `struct ƛint_int' has to be declared before the function whose
    ;; signature first mentioned it, so the forms that function produced
    ;; are lifted off and the structs slid underneath them
    (let* ((processed (match-sex-form (car forms) acc))
           (new (take-until processed acc)))
      (process-rec (cdr forms)
                   (append new (flush-closure-structs!) acc))))))

(define (take-until forms tail)
  (if (eq? forms tail)
      (list)
      (cons (car forms) (take-until (cdr forms) tail))))

(define (macroexpand macro-form rest-forms)
  ;; `(defmacro (two) 2)' expands to `2', and `($ (fn a ...) (fn b ...))'
  ;; to two forms, spliced where the macro was written. `($)' expands to
  ;; nothing.
  ;;
  ;; Everything else is one form, a list whose head is itself a form
  ;; included: `((make-adder 10) 5)' calls what `make-adder' returned,
  ;; and without `$' to mark a splice there is no telling that from a
  ;; list of the two forms `(make-adder 10)' and `5'.
  (let ((res (apply-macro macro-form))
        (src (form-source macro-form)))
    ;; An expansion is fresh structure with no location of its own. Give
    ;; it the call site's, the way cpp attributes a macro body to where
    ;; the macro was used
    (cond
     ((splice-form? res)
      (append (map (lambda (f) (stamp-form-source! f src)) (cdr res))
              rest-forms))
     ((null? res) rest-forms)
     (else
      (cons (stamp-form-source! res src) rest-forms)))))

(define (splice-form? form)
  (and (pair? form) (list? form) (eq? '$ (car form))))

(define (match-sex-form sex-form acc)
  (match sex-form
    ((or ('fn . _)
         ('pub 'fn . _)
         ('extern 'fn . _)) (process-fn sex-form acc))
    ((or ('struct . _)
         ('pub 'struct . _)) (process-struct sex-form acc))
    ((or ('union . _)
         ('pub 'union . _)) (process-struct sex-form acc))
    ((or ('enum . _)
         ('pub 'enum . _)) (process-struct sex-form acc))
    ((or ('var . _)
         ('pub 'var . _)
         ('extern 'var . _)) (process-global-var sex-form acc))
    (('include . includes) (process-includes sex-form includes acc))
    ((or ('define name . _)
         ('pub 'define name . _))
     (add-define name sex-form)
     (cons sex-form acc))
    (('comment . _) (cons sex-form acc))

    (('import . modules)
     (process-imports (get-modules-public-forms modules) acc))

    ((or ('defmacro . rest)
         ('pub 'defmacro . rest)) (defmacro rest) acc)

    ((or ('typedef new-type target)
         ('pub 'typedef new-type target))
     (add-typedef new-type sex-form)
     (process-typedef sex-form new-type target acc))

    (else (sex-error sex-form "unknown top level form" sex-form))))

(define (process-includes sex-form includes acc)
  ;; consume (include ...) form and add to acc
  ;; (include <inc>) for each include
  (let process ((includes includes)
                (acc acc))
    (if (null? includes)
        acc
        (process (cdr includes)
                 (cons (copy-form-source! sex-form `(include ,(car includes)))
                       acc)))))

(define (process-imports module-public-forms acc)
  ;; consume (import ...) form and process imports so
  ;; data types end up in types db
  (fold match-sex-form acc module-public-forms))

(define (macro-expand form)
  "Walk the form recursively and expand all macros, until none is left."
  (walk-form
   form
   (lambda (subform env)
     (if (macro? subform)
         (cons walk-embed-result (macroexpand subform (list)))
         subform))
   #f))

;;; walk-form and friends: form walker with various abilities.
;;; By default, replaces walked form with walk-fn result But may
;;; perform additional operations depending of what the walk function
;;; has requested.

;;; For inspiration, see SBCL's walk.lisp and their template
;;; system.

(define walk-embed-result (gensym)
  ;; For cases when result is a list which must be embedded in the
  ;; form, e.g. when it returned from a macro
  )

(define (walk-form form walk-fn env)
  (if (atom? form) form
      (let ((new-form (walk-fn form env)))
        (cond ((not (eq? form new-form))
               (walk-form new-form walk-fn env))
              (else
               (let ((new-car (walk-form (car new-form) walk-fn env))
                     (new-cdr (walk-form (cdr new-form) walk-fn env)))
                 (cond ((and (pair? new-car)
                             (eq? (car new-car) walk-embed-result))
                        (append (cdr new-car) new-cdr))
                       (else
                        (recons new-form new-car new-cdr)))))))))

;;; Typdef

(define (process-typedef form new-type target acc)
  (cons (copy-form-source! form `(typedef ,target ,new-type)) acc))

;;; Fn processing
;;;
;;; A string as the first body form is a docstring. In the generated
;;; C code it will be placed as a C commentary just before the function
;;; definition (actually that works for all blocky things: enum, struct, union as well).

(define (fn-header-length fn-form)
  (if (memq (first fn-form) '(pub extern)) 5 4))

(define (fn-core form)
  ;; The (fn name args rettype . body) list, without pub/extern
  (if (memq (first form) '(pub extern))
      (cdr form)
      form))

(define (take-leading-docstring forms)
  ;; If FORMS starts with a string, possibly after comment forms, return
  ;; that string and FORMS without it. Otherwise #f and FORMS unchanged
  (let loop ((fs forms) (prefix (list)))
    (match fs
      (() (values #f forms))
      (((and cmt ('comment . _)) . rest)
       (loop rest (cons cmt prefix)))
      (((? string? doc) . rest)
       (values doc (append (reverse prefix) rest)))
      (_ (values #f forms)))))

(define (extract-fn-docstring fn-form)
  (let ((lift
         (lambda (proto body)
           (let-values (((doc rest) (take-leading-docstring body)))
             (if doc
                 (values doc (copy-form-source! fn-form (append proto rest)))
                 (values #f fn-form))))))
    (match fn-form
      (('pub 'fn name args ret . body)
       (lift `(pub fn ,name ,args ,ret) body))
      (('extern 'fn name args ret . body)
       (lift `(extern fn ,name ,args ,ret) body))
      (('fn name args ret . body)
       (lift `(fn ,name ,args ,ret) body))
      (_ (values #f fn-form)))))

(define (extract-aggregate-docstring form)
  ;; A string immediately after the name is the docstring; comments
  ;; between name and fields are not skipped, they already confuse the
  ;; writer
  (match form
    (('pub (and kind (or 'struct 'union 'enum))
           (? symbol? name) (? string? doc) . rest)
     (values doc (copy-form-source! form `(pub ,kind ,name ,@rest))))
    (((and kind (or 'struct 'union 'enum))
      (? symbol? name) (? string? doc) . rest)
     (values doc (copy-form-source! form `(,kind ,name ,@rest))))
    (_ (values #f form))))

(define (with-docstring doc form acc)
  ;; acc is newest-first; FORM is consed last so the final reverse
  ;; emits the comment immediately before the declaration
  (cons form
        (if doc
            (cons (list 'comment doc) acc)
            acc)))

(define (strip-fn-header-comments fn-form)
  ;; ([pub|extern] fn name arglist rettype). Comments in the body are
  ;; left in place as ordinary statements and preserved into the
  ;; generated C.
  (strip-header-comments fn-form (fn-header-length fn-form)))

(define (process-fn sex-fn-raw acc)
  (let-values (((doc sex-fn)
                (extract-fn-docstring (strip-fn-header-comments sex-fn-raw))))
    (let* ((boundary (fn-header-length sex-fn))
           ;; only the header is expanded up front. A macro in the body
           ;; is expanded during the walk, where `(type-of x)' can still
           ;; be answered from the scope it was written in
           (expanded (append (macro-expand (take sex-fn boundary))
                             (drop sex-fn boundary))))
      (add-name-type! (sex-fn-name expanded) (fn-type-of expanded))
      (let* ((env       (make-fn-env expanded))
             ;; the parameters are the body's outermost scope
             (header    (map resolve-closure-types (take expanded boundary)))
             (body      (walk-body (drop expanded boundary) env))
             (processed (append header body)))
        (with-docstring doc processed
                        (append (hash-table-ref env :lambda-aux-code) acc))))))

(define (make-fn-env fn-form)
  (let ((env (make-hash-table))
        (parameters (make-hash-table)))
    (for-each (lambda (param)
                (when (and (pair? param) (pair? (cdr param)))
                  (hash-table-set! parameters (first param) (second param))))
              (sex-fn-arglist fn-form))
    (set! (hash-table-ref env :fn-name) (sex-fn-name fn-form))
    (set! (hash-table-ref env :lambda-counter) 0)
    (set! (hash-table-ref env :lambda-aux-code) (list))
    (set! (hash-table-ref env :returns) (sex-fn-return-type fn-form))
    (set! (hash-table-ref env :scopes) (list parameters))
    env))

;;; (fn sum ((a int) (b int)) int ...) -> (fn ((int) (int)) int).
;;; Anything that is not a plain (name type) -- a variadic tail -- goes
;;; through untouched.
(define (fn-type-of fn-form)
  `(fn ,(map (lambda (param)
               (if (and (pair? param) (= 2 (length param)) (named-arg? param))
                   (list (second param))
                   param))
             (sex-fn-arglist fn-form))
       ,(sex-fn-return-type fn-form)))

(define (aux-name! env make)
  (let ((counter (hash-table-ref env :lambda-counter)))
    (set! (hash-table-ref env :lambda-counter) (+ counter 1))
    (make (hash-table-ref env :fn-name) counter)))

(define (add-aux-code! env forms)
  (set! (hash-table-ref env :lambda-aux-code)
    (append forms (hash-table-ref env :lambda-aux-code))))


;;; The scope chain
;;;
;;; `(do (var c int 9) ...)' declares a `c' that ends with the block, so
;;; a closure-typed `c' outside it is still a closure after it. `do' and
;;; `for' each open a frame; innermost first.

(define (declare-name! env name type)
  (hash-table-set! (car (hash-table-ref env :scopes)) name type))

(define (lookup-name env name)
  (let search ((scopes (hash-table-ref env :scopes)))
    (and (pair? scopes)
         (or (hash-table-ref/default (car scopes) name #f)
             (search (cdr scopes))))))

(define (with-scope env body)
  (let ((enclosing (hash-table-ref env :scopes)))
    (set! (hash-table-ref env :scopes) (cons (make-hash-table) enclosing))
    (let ((walked (body)))
      (set! (hash-table-ref env :scopes) enclosing)
      walked)))

;;; The body walk
;;;
;;; One pass in statement order: lifts lambdas and closures out,
;;; resolves closure types to the struct that stands for them, records
;;; what each declaration binds, and rewrites a call whose head is a
;;; closure.
;;;
;;; Not `walk-form': it has no event for leaving a scope, and it hands
;;; the walk function every cdr-tail as well, so the `f' in `(g f)'
;;; arrives as `(f)' and reads as a call of its own.

(define (walk-body forms env)
  (append-map (lambda (form)
                (let ((walked (walk-statement form env)))
                  (if (and (pair? walked) (eq? (car walked) walk-embed-result))
                      (cdr walked)
                      (list walked))))
              forms))

(define (walk-statement form env)
  (cond
   ((not (list? form)) form)
   ((null? form) form)
   ;; expanded here rather than before the walk, so the macro body can
   ;; ask `(type-of x)' about a local the walk has already passed
   ((macro? form)
    (let ((expansion (parameterize
                         ((current-type-of
                           (lambda (queried)
                             (unresolve-closure-types
                              (expression-type queried env)))))
                       (macroexpand form (list)))))
      (if (and (pair? expansion) (null? (cdr expansion)))
          (walk-statement (car expansion) env)
          (cons walk-embed-result (walk-body expansion env)))))
   (else
    (case (car form)
      ((do for) (with-scope env (lambda () (walk-parts form env))))

      ((lambda)
       (let ((name (aux-name! env make-lambda-name)))
         (add-aux-code! env (lift-lambda name form))
         name))

      ((closure)
       (if (closure-type? form)
           ;; a type, not an expression: the struct for its signature
           (copy-form-source! form `(struct ,(register-closure-type! form form)))
           (let ((base (aux-name! env make-closure-name)))
             (let-values (((construct lifted) (lift-closure base form env)))
               (add-aux-code! env (fold match-sex-form (list) lifted))
               (copy-form-source!
                form
                `(,construct ,@(map (lambda (capture)
                                      (walk-statement (capture-argument capture)
                                                      env))
                                    (fourth form))))))))

      ((var)
       ;; the initializer is walked before the name it binds is in scope
       (let* ((walked (resolve-wildcard (walk-parts form env) env))
              (bound  (if (>= (length walked) 4)
                          (copy-form-source!
                           walked
                           (append (take walked 3)
                                   (cons (convert-to-closure (third walked)
                                                             (fourth walked)
                                                             env walked)
                                         (drop walked 4))))
                          walked)))
         (when (>= (length bound) 3)
           (declare-name! env (second bound) (third bound)))
         bound))

      ((return)
       (let ((walked (walk-parts form env)))
         (if (>= (length walked) 2)
             (copy-form-source!
              walked
              (cons 'return
                    (cons (convert-to-closure (hash-table-ref env :returns)
                                              (second walked) env walked)
                          (cddr walked))))
             walked)))

      (else
       (let ((closure (receiver-closure-type (car form) env)))
         (if closure
             (copy-form-source!
              form
              `(,(register-closure-call! closure form)
                ,(walk-statement (car form) env)
                ,@(map (lambda (argument) (walk-statement argument env))
                       (cdr form))))
             (convert-arguments (walk-parts form env) env))))))))

;;; `(var n _ (strlen s))' becomes `(var n size-t (strlen s))'.
(define (resolve-wildcard form env)
  (if (and (>= (length form) 3) (wildcard-type? (third form)))
      (let ((declared (parse-type (third form))))
        (solve-wildcards! declared
                          (and (>= (length form) 4) (fourth form))
                          form env)
        (let ((written (unparse-type declared)))
          (cond
           ;; `?' is what a name from an unparsed header types as, and
           ;; the writer has no spelling for it
           ((mentions? written '?)
            (sex-error form "type of this is unknown; write it out"
                       (second form)))
           ((mentions? written '_)
            (sex-error form "cannot infer the type of" (second form)))
           (else
            (copy-form-source! form
                               (cons (first form)
                                     (cons (second form)
                                           (cons written (cdddr form)))))))))
      form))

;;; `(* _)' against `(* (struct point))' solves only the wildcard inside
;;; the pointer; a bare `_' is the same with nothing around it.
;;;
;;; `#(0 1 4 9)' has no type of its own, so it cannot answer a bare `_',
;;; but its elements still solve the hole in `(¤ _ 4)' -- each one is
;;; unified with the element type, which is what makes `(¤ _ 4)' worth
;;; writing at all.
(define (solve-wildcards! declared initializer form env)
  (cond
   ((not initializer)
    (sex-error form "cannot infer the type of" (second form)))
   ((brace-initializer? initializer)
    (let ((element (and (array-type? declared) (array-elt declared))))
      (unless element
        (sex-error form "cannot infer the type of" (second form)))
      (for-each (lambda (written)
                  (let ((type (expression-type written env)))
                    (when type
                      (unify element
                             (parse-type (resolve-closure-types type))
                             form))))
                (vector->list initializer))))
   (else
    (let ((type (expression-type initializer env)))
      (unless type
        (sex-error form "cannot infer the type of" (second form)))
      ;; `(closure ((int)) int)' is spelled `(struct ƛint_int)'
      ;; everywhere past this point, and is what `parse-type' knows
      (unify declared (parse-type (resolve-closure-types type)) form)))))

;;; `#(0 1 4 9)', as against the compound literal `#(T : ...)'
(define (brace-initializer? form)
  (and (vector? form)
       (null? (cdr (list-split (vector->list form) ':)))))

(define (wildcard-type? type) (mentions? type '_))

(define (mentions? type word)
  (cond ((eq? type word) #t)
        ((list? type) (any (lambda (part) (mentions? part word)) type))
        (else #f)))

;;; `(each xs compare)' where `each' takes a closure: the argument is
;;; checked against the parameter that signature wrote.
(define (convert-arguments form env)
  (let ((signature (and (symbol? (car form)) (get-name-type (car form)))))
    (if (and (list? signature) (= 3 (length signature)) (eq? 'fn (car signature)))
        (copy-form-source!
         form
         (cons (car form)
               (map (lambda (argument expected)
                      (if expected
                          (convert-to-closure expected argument env form)
                          argument))
                    (cdr form)
                    (parameter-types signature (length (cdr form))))))
        form)))

;;; One per argument, #f past the end of the parameter list -- a
;;; variadic tail has nothing written to check against
(define (parameter-types signature count)
  (let pair ((params (second signature)) (remaining count) (acc (list)))
    (cond
     ((zero? remaining) (reverse acc))
     ((null? params) (pair params (- remaining 1) (cons #f acc)))
     (else (pair (cdr params) (- remaining 1)
                 (cons (unwrap-type (car params)) acc))))))

(define (walk-parts form env)
  (copy-form-source! form (walk-body form env)))


(define (as-closure-type type)
  (cond
   ((closure-type? type) type)
   ((and (list? type) (= 2 (length type)) (eq? 'struct (car type)))
    (hash-table-ref/default +closure-structs+ (second type) #f))
   (else #f)))

(define (receiver-closure-type expr env)
  (as-closure-type (expression-type expr env)))

;;; The type of an lvalue path, from known declarations -- a name, and
;;; what can be reached from one by subscripting, dereferencing and
;;; member access.

;;; The type of an expression as a surface type, or #f when nothing
;;; here can say. Every form it visits is recorded in `form-type', so
;;; asking once types the whole subtree.
;;;
;;;   42          -> int              (. p x)    -> that field's type
;;;   "hi"        -> (* const char)   (& p)      -> (* (struct point))
;;;   (area 3 4)  -> what `area' returns
(define (expression-type expr env)
  (set-form-type! expr (compute-expression-type expr env)))

(define (compute-expression-type expr env)
  (cond
   ((and (number? expr) (exact? expr)) 'int)
   ((number? expr) 'double)
   ((string? expr) '(* const char))
   ((char? expr) 'char)
   ((memq expr '(true false)) 'bool)
   ;; `#(T : ...)' carries its own type; a bare `#(...)' has none and
   ;; takes one from whatever it is being written into
   ((vector? expr)
    (let ((parts (list-split (vector->list expr) ':)))
      (and (pair? (cdr parts))
           (let ((written (car parts)))
             (if (= 1 (length written)) (car written) written)))))
   ((symbol? expr) (or (lookup-name env expr) (get-name-type expr)))
   ((not (and (list? expr) (pair? expr))) #f)
   ;; a call of no arguments is still a call
   ((< (length expr) 2)
    (and (symbol? (car expr)) (get-return-type (car expr))))
   (else
    (case (car expr)
      ;; subscripting an array gives its element type, and a pointer
      ;; subscripts the same way
      ((¤) (let ((base (expression-type (second expr) env)))
             (or (array-element-type base) (pointer-target base))))
      ((&) (and (= 2 (length expr))
                (let ((target (expression-type (second expr) env)))
                  (and target `(* ,target)))))
      ;; unary `*' is a dereference; with two operands it is a product
      ((*) (if (= 2 (length expr))
               (pointer-target (expression-type (second expr) env))
               (arithmetic-type expr env)))
      ((dot-access) (member-path-type (expression-type (second expr) env)
                                      (cddr expr)))
      ((->) (member-path-type (pointer-target (expression-type (second expr) env))
                              (cddr expr)))
      ((cast) (and (= 3 (length expr)) (third expr)))
      ((sizeof) 'size-t)
      ((== != < > <= >= c-and c-or !) 'bool)
      ((+ - / %) (arithmetic-type expr env))
      ;; otherwise a call: a closure answers with its own return type,
      ;; anything else with what its signature says
      (else
       (let ((closure (receiver-closure-type (car expr) env)))
         (if closure
             (third closure)
             (and (symbol? (car expr)) (get-return-type (car expr))))))))))

;;; C's usual arithmetic conversions, far enough to answer `_':
;;;
;;;   (+ i d)     i int, d double  -> double
;;;   (+ i l)     l long           -> long
;;;   (+ g d)     g float          -> double
;;;   (+ (& p) 1)                  -> (* (struct point))
;;;
;;; `unsigned int' against `int' answers `int', where C says otherwise.
(define (arithmetic-type expr env)
  (fold (lambda (operand joined)
          (arith-join joined (expression-type operand env)))
        #f
        (cdr expr)))

(define (arith-join left right)
  (cond
   ((not left) right)
   ((not right) left)
   ((equal? left right) left)
   (else
    (let ((l (underlying (parse-type left)))
          (r (underlying (parse-type right))))
      (cond
       ((or (ptr-type? l) (array-type? l)) left)
       ((or (ptr-type? r) (array-type? r)) right)
       ((< (conversion-rank l) (conversion-rank r)) right)
       (else left))))))

;;; `char' and `short' promote to `int', so the ranks start there
(define (conversion-rank type)
  (let ((words (and (prim-type? type) (prim-name type))))
    (cond
     ((not words) 1)
     ((memq 'double words) (if (memq 'long words) 7 6))
     ((memq 'float words) 5)
     ((memq 'long words) (if (= 2 (count (lambda (w) (eq? w 'long)) words)) 4 3))
     (else 1))))


;;; `(* const char)' is flat: everything after the `*' is the target.
;;; `(& x)' builds the nested `(* (struct point))' and `unparse-type'
;;; writes the flat `(* struct point)', so both spellings turn up.
(define (pointer-target type)
  (and (list? type)
       (>= (length type) 2)
       (eq? '* (car type))
       (unwrap-type (cdr type))))

(define (member-path-type type fields)
  (if (null? fields)
      type
      (member-path-type (field-type type (car fields)) (cdr fields))))

(define (field-type type field)
  (let ((name (cond ((symbol? type) type)
                    ((and (list? type)
                          (= 2 (length type))
                          (memq (car type) '(struct union)))
                     (second type))
                    (else #f))))
    (and name
         (let* ((fields (get-fields name))
                (entry (and fields (assq field fields))))
           (and entry (second entry))))))

(define (make-lambda-name enclosing-fn-name counter)
  (string->symbol
   (fmt #f "λ" counter "_" enclosing-fn-name)))

(define (lift-lambda name form)
  (match form
    (('lambda arglist ret-type . body)
     (process-fn (copy-form-source! form `(fn ,name ,arglist ,ret-type ,@body))
                 (list)))
    (else (sex-error form "malformed lambda" form))))

;;; Closures
;;;
;;; A closure is a function pointer and an inline environment, so the
;;; value owns its captures and nothing is allocated. The type
;;; `(closure ((int)) int)' becomes one struct per signature, shared by
;;; every closure with that signature. The captures live in `env' as a record
;;; only the lifted body knows the shape of, which is why `env' is
;;; max_align_t rather than char -- it has to be aligned for whatever
;;; ends up in it.
;;;
;;; The expression becomes three hoisted definitions -- the capture
;;; struct, the lifted body, and a constructor -- and is replaced by a
;;; call to the constructor, so the captures are evaluated as ordinary
;;; arguments at the point the closure is written.

;;; How much of a closure is environment, in bytes. Counted in bytes
;;; so a closure that fits where it was written will also fit
;;; elsewhere. Captures are only by value, and never allocated on
;;; heap. Anything that is more than 16 bytes should be stored as a
;;; pointer, and memory management is entirely up to caller

(define +closure-env-bytes+ 16)

;;; +closure-env-bytes+ for maximum capacity, max-align-t for
;;; effectiveness, hence union
(define +closure-env-type+ 'ƛenv)

(define (closure-env-declaration)
  `(union ,+closure-env-type+ ((align max-align-t)
                               (bytes (¤ char ,+closure-env-bytes+)))))

(define +closure-structs+ (make-hash-table))
(define +closure-forwards+ (make-hash-table))
(define *pending-closure-structs* (list))

(define (closure-expression? form)
  (and (list? form)
       (pair? form)
       (eq? 'closure (car form))
       (>= (length form) 4)))

(define (closure-type? form)
  (and (pair? form)
       (eq? 'closure (car form))
       (= 3 (length form))))

;;; A type spelling becomes an identifier deterministically, so two
;;; translation units have the same signatures for the same closure
;;; types: (* const char) -> p_const_char, (closure ((int)) int) ->
;;; closure_int_int.
(define (mangle-type type)
  (cond
   ((symbol? type) (mangle-word (symbol->string type)))
   ((number? type) (number->string type))
   ((null? type) "void")
   ((pair? type) (string-intersperse (map mangle-type (mangle-head type)) "_"))
   (else (sex-error type "cannot mangle type" type))))

(define (mangle-head type)
  (case (car type)
    ((*) (cons 'p (cdr type)))
    ((¤) (cons 'a (cdr type)))
    (else type)))

(define (mangle-word word)
  (list->string
   (map (lambda (c)
          (if (or (char-alphabetic? c) (char-numeric? c)) c #\_))
        (string->list word))))

(define (aggregates-in type)
  (cond
   ((not (list? type)) (list))
   ((and (= 2 (length type))
         (memq (car type) '(struct union))
         (symbol? (second type)))
    (list type))
   (else (append-map aggregates-in type))))

;;; Extract aggregate types from the closure's signature to forward
;;; declare them before the closure, so they can be referenced in the
;;; closure. Particularly useful for complex cases like fixed point
;;; combinator, etc.
(define (forward-declare-aggregates! type src-form)
  (for-each
   (lambda (aggregate)
     (let ((name (second aggregate)))
       (unless (or (hash-table-exists? +closure-forwards+ name)
                   (get-tag-info name))
         (hash-table-set! +closure-forwards+ name #t)
         (set! *pending-closure-structs*
           (cons (copy-form-source! src-form aggregate)
                 *pending-closure-structs*)))))
   (delete-duplicates (aggregates-in type))))

(define (closure-struct-name type)
  ;; the glyph says `closure' already, so the tag is just the signature
  (string->symbol (string-append "ƛ"
                                 (mangle-type (second type))
                                 "_"
                                 (mangle-type (third type)))))

;;; The code pointer takes the environment first; everything else is
;;; the closure's own signature.
(define (closure-code-type type)
  `(fn (((* void)) ,@(second type)) ,(third type)))

;;; Emitted once per signature, before the toplevel form that first
;;; needed it.
(define (register-closure-type! type src-form)
  (let ((name (closure-struct-name type)))
    (unless (hash-table-exists? +closure-structs+ name)
      (forward-declare-aggregates! type src-form)
      (hash-table-set! +closure-structs+ name type)
      (let ((form (copy-form-source!
                   src-form
                   `(struct ,name ((code ,(closure-code-type type))
                                   (env (union ,+closure-env-type+)))))))
        (register-aggregate! form)
        (set! *pending-closure-structs*
          (cons form *pending-closure-structs*))))
    name))

;;; Every closure type in FORM becomes the struct for its signature,
;;; registering it on the way. The walker does this for function bodies
;;; and headers; globals come through here instead.
;;; The inverse of `resolve-closure-types', for what a macro is shown:
;;; `(struct ƛint_int)' is a generated name, and `(closure ((int)) int)'
;;; is what was written and what a `type-match' pattern says.
(define (unresolve-closure-types type)
  (or (as-closure-type type)
      (if (list? type)
          (map unresolve-closure-types type)
          type)))

(define (resolve-closure-types form)
  (cond
   ;; `(closure args ret captures . body)' is an expression, and one is
   ;; lifted into the function it was written in. At toplevel, or in a
   ;; struct field, there is none
   ((closure-expression? form)
    (sex-error form "a closure can only be written inside a function" form))
   ((closure-type? form)
    (copy-form-source! form `(struct ,(register-closure-type! form form))))
   ((list? form)
    (copy-form-source! form (map resolve-closure-types form)))
   (else form)))

(define +closure-conversions+ (make-hash-table))

;;; `(var c (closure ((int)) int) sum)' becomes `ƛint_int_fromfn(sum)'.
;;; The function pointer goes in the environment and a thunk reads it
;;; back out, so one thunk serves every function of that signature:
;;;
;;;   struct ƛint_int_fnptr { int (*f)(int); };
;;;   static int ƛint_int_fnthunk (void *ƛe, int ƛa0) {
;;;       struct ƛint_int_fnptr *ƛcaptures = ƛe;
;;;       return ƛcaptures->f(ƛa0);
;;;   }
(define (register-fn-conversion! type src-form)
  (let* ((closure (closure-struct-name type))
         (convert (suffixed closure "_fromfn"))
         (record  (suffixed closure "_fnptr"))
         (thunk   (suffixed closure "_fnthunk"))
         (returns (third type))
         (pointer `(fn ,(second type) ,returns))
         (params  (map (lambda (argument index)
                         (list (string->symbol (fmt #f "ƛa" index))
                               (unwrap-type argument)))
                       (second type)
                       (iota (length (second type))))))
    (unless (hash-table-exists? +closure-conversions+ convert)
      (hash-table-set! +closure-conversions+ convert #t)
      (let ((call `((-> ƛcaptures f) ,@(map first params))))
        (for-each
         (lambda (emitted)
           (set! *pending-closure-structs*
             (cons (copy-form-source! src-form emitted)
                   *pending-closure-structs*)))
         (list
          `(struct ,record ((f ,pointer)))

          `(fn ,thunk ((ƛe (* void)) ,@params) ,returns
               (var ƛcaptures (* (struct ,record)) ƛe)
               ,(if (eq? 'void returns) call `(return ,call)))

          `(fn ,convert ((f ,pointer)) (struct ,closure)
               (var ƛc (struct ,closure))
               (= (dot-access ƛc code) ,thunk)
               (var ƛcaptures (* (struct ,record))
                    (cast (& (dot-access ƛc env)) (* (struct ,record))))
               (= (-> ƛcaptures f) f)
               (return ƛc))))))
    convert))

;;; A bare function is a closure that captures nothing, so it converts
;;; wherever one is expected. The reverse cannot: a closure has an
;;; environment and a function pointer has nowhere to put it.
(define (convert-to-closure expected value env form)
  (let ((closure (as-closure-type expected))
        (actual  (expression-type value env)))
    (if (and closure
             (list? actual)
             (= 3 (length actual))
             (eq? 'fn (car actual))
             (equal? (cdr actual) (cdr closure)))
        (copy-form-source! form
                           (list (register-fn-conversion! closure form) value))
        value)))

(define +closure-calls+ (make-hash-table))

;;; The helper a closure call is routed through: `(f 1)' becomes
;;; `ƛint_int_call(f, 1)', which unpacks the receiver into
;;; `ƛc.code(&ƛc.env, ƛa0)' inside. The receiver arrives as an argument,
;;; so it is evaluated once: unpacked at the call site instead,
;;; `([table (++ i)] 10)' would read
;;; `table[++i].code(&table[++i].env, 10)' and bump `i' twice.
(define (register-closure-call! type src-form)
  (let* ((closure (closure-struct-name type))
         (helper  (suffixed closure "_call"))
         (returns (third type))
         (params  (map (lambda (arg index)
                         (list (string->symbol (fmt #f "ƛa" index))
                               (unwrap-type arg)))
                       (second type)
                       (iota (length (second type))))))
    (unless (hash-table-ref/default +closure-calls+ helper #f)
      (hash-table-set! +closure-calls+ helper #t)
      (let ((call `((dot-access ƛc code)
                    (& (dot-access ƛc env))
                    ,@(map first params))))
        (set! *pending-closure-structs*
          (cons (copy-form-source!
                 src-form
                 `(fn ,helper ((ƛc (struct ,closure)) ,@params) ,returns
                      ,(if (eq? 'void returns) call `(return ,call))))
                *pending-closure-structs*))))
    helper))

;;; An argument type is written wrapped: `(int)' in `((int) (float))'
(define (unwrap-type type)
  (if (and (list? type) (= 1 (length type)))
      (car type)
      type))

(define (flush-closure-structs!)
  (let ((pending *pending-closure-structs*))
    (set! *pending-closure-structs* (list))
    pending))

;;; Lowering

(define (make-closure-name enclosing-fn-name counter)
  (string->symbol (fmt #f "ƛ" counter "_" enclosing-fn-name)))

(define (suffixed name suffix)
  (string->symbol (string-append (symbol->string name) suffix)))

;;; `(closure ((b int)) int (n) ...)' captures `n' under its own name
;;; and takes its type from wherever it was declared. `((pa (& a)))'
;;; names the capture and gives the expression it holds, so `pa' is a
;;; `(* int)' inside the body and `a' is never mentioned there.
(define (capture-binding capture form env)
  (let ((name (capture-name capture)))
    (unless (symbol? name)
      (sex-error form "a closure capture needs a name" capture))
    (let ((type (if (pair? capture)
                    (expression-type (capture-argument capture) env)
                    (lookup-name env name))))
      (unless type
        (sex-error form "cannot infer what is captured as" name))
      (list name type))))

(define (capture-name capture)
  (if (pair? capture) (first capture) capture))

;;; What the constructor is handed: the name itself, or the expression
;;; written beside it
(define (capture-argument capture)
  (if (pair? capture) (second capture) capture))

(define (lift-closure base form env)
  (match form
    (('closure arglist ret-type captures . body)
     (let* ((caps      (map (lambda (c) (capture-binding c form env)) captures))
            (record    (suffixed base "_captures"))
            (code      (suffixed base "_code"))
            (construct (suffixed base "_make"))
            (type      `(closure ,(map (lambda (p) (list (second p))) arglist)
                                 ,ret-type))
            (closure   (register-closure-type! type form)))
       (values
        construct
        (map (lambda (f) (copy-form-source! form f))
             ;; C has no empty struct, and a closure over nothing needs
             ;; no record to point at
             (append
              (if (null? caps)
                  (list)
                  (list `(struct ,record ,caps)))

              (list
               `(fn ,code ((ƛe (* void)) ,@arglist) ,ret-type
                    ,@(if (null? caps)
                          (list)
                          `((var ƛcaptures (* (struct ,record)) ƛe)
                            ,@(map (lambda (cap)
                                     `(var ,(first cap) ,(second cap)
                                           (-> ƛcaptures ,(first cap))))
                                   caps)))
                    ,@body)

               `(fn ,construct ,caps (struct ,closure)
                    (var ƛc (struct ,closure))
                    ,@(if (null? caps)
                          (list)
                          `((static-assert
                             (<= (sizeof (struct ,record))
                                 (sizeof (dot-access ƛc env)))
                             "closure captures do not fit the inline environment")))
                    (= (dot-access ƛc code) ,code)
                    ,@(if (null? caps)
                          (list)
                          `((var ƛcaptures (* (struct ,record))
                                 (cast (& (dot-access ƛc env)) (* (struct ,record))))
                            ,@(map (lambda (cap)
                                     `(= (-> ƛcaptures ,(first cap)) ,(first cap)))
                                   caps)))
                    (return ƛc))))))))
    (else (sex-error form "malformed closure" form))))

;;; Structs

;;; Record the named structs, unions and enums in the type database
(define (process-struct sex-struct acc)
  (let-values (((doc form) (extract-aggregate-docstring sex-struct)))
    ;; Resolve before registering: a field of closure type has to reach
    ;; the type database as the struct it becomes, or member access
    ;; through it finds nothing
    (let ((form (resolve-closure-types form)))
      (register-aggregate! form)
      (with-docstring doc form acc))))

(define (register-aggregate! form)
  (let* ((f (if (eq? (car form) 'pub) (cdr form) form))
         (name (and (pair? (cdr f)) (symbol? (cadr f)) (cadr f))))
    ;; An anonymous aggregate has a field list where the name would be,
    ;; and nothing can refer to it by name anyway
    (when name
      (case (car f)
        ((struct) (add-struct name form))
        ((union) (add-union name form))
        ((enum) (add-enum name form))))))

(define (process-global-var sex-var acc)
  ;; A global is not walked for lambdas, but its type still has to stop
  ;; saying `closure' before the writer sees it
  (let* ((form (resolve-closure-types sex-var))
         (core (if (memq (car form) '(pub extern)) (cdr form) form)))
    (when (and (pair? (cdr core)) (pair? (cddr core)) (symbol? (second core)))
      (add-name-type! (second core) (third core)))
    (cons form acc)))

;;; Utils
(define (non-empty-list? form)
  (and (list? form)
       (not (null? form))))

(define (sex-fn? form)
  "The `form` must be toplevel.
Returns #f if the form is not a function, returns the form otherwise"
  (match form
    ((or ('fn . _)
         ('pub 'fn . _)
         ('extern 'fn . _)) form)
    (else #f)))

(define (sex-fn-public? fn-form)
  (eq? (first fn-form) 'pub))

(define (sex-fn-name fn-form)
  (assert (sex-fn? fn-form)
          (fmt #f "Form " fn-form " is not a function"))
  (second (fn-core fn-form)))

(define (sex-fn-arglist fn-form)
  (assert (sex-fn? fn-form)
          (fmt #f "Form " fn-form " is not a function"))
  (third (fn-core fn-form)))

(define (sex-fn-return-type fn-form)
  (assert (sex-fn? fn-form)
          (fmt #f "Form " fn-form " is not a function"))
  (fourth (fn-core fn-form)))

(define (sex-fn-prototype fn-form)
  "Returns all except body"
  (assert (sex-fn? fn-form)
          (fmt #f "Form " fn-form " is not a function"))
  (take fn-form (fn-header-length fn-form)))

(define (sex-fn-body fn-form)
  (assert (sex-fn? fn-form)
          (fmt #f "Form " fn-form " is not a function"))
  (drop fn-form (fn-header-length fn-form)))