;;; Sex semantic engine

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

(export/rename (process semen-process))

;;; 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
    (process-rec (cdr forms)
                 (match-sex-form (car forms) acc)))))

(define (macroexpand macro-form rest-forms)
  ;; We want to replace macro with its expansion. The problem is,
  ;; top-level macro can return either a single form, or a list of
  ;; forms, when it for example generates some aux
  ;; structures/functions/typedefs.
  ;;
  ;; Single form we just cons to the top of rest-forms, but multiple
  ;; forms have to be appended to the rest-forms.
  (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
    (if (list? (car res))
        (append (map (lambda (f) (stamp-form-source! f src)) res)
                rest-forms)
        (cons (stamp-form-source! res src) rest-forms))))

(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 _) (cons sex-form 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-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* ((expanded (macro-expand sex-fn))
           (env (make-hash-table))
           (processed
            (walk-form
             expanded
             fn-walker
             (begin
               (set! (hash-table-ref env :fn-name) (sex-fn-name expanded))
               (set! (hash-table-ref env :lambda-counter) 0)
               (set! (hash-table-ref env :lambda-aux-code) (list))
               env))))
      (with-docstring doc processed
                      (append (hash-table-ref env :lambda-aux-code) acc)))))

(define (fn-walker form env)
  (if (eq? 'lambda (car form))
      (let ((lambda-name (make-lambda-name (hash-table-ref env :fn-name)
                                           (hash-table-ref env :lambda-counter))))
        (set! (hash-table-ref env :lambda-aux-code)
          (append (make-aux-lambda-struct lambda-name form)
                  (hash-table-ref env :lambda-aux-code)))
        (set! (hash-table-ref env :lambda-counter)
          (+ (hash-table-ref env :lambda-counter) 1))
        lambda-name)
      form))

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

(define (make-aux-lambda-struct name form)
  (match form
    (('lambda ret-type arglist captures . body)
     ;; Captures are ignored for now, but
     ;; we'll need them for TODO: closures support
     (process-fn (copy-form-source! form `(fn ,ret-type ,name ,arglist ,@body))
                 (list)))
    (else (sex-error form "malformed lambda" 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)))
    (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)
  (cons sex-var 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)))