;; Desugar the Haskell surface AST into a smaller core AST.
;;
;; Eliminates the three surface-only shapes produced by the parser:
;;   :where BODY DECLS     →  :let DECLS BODY
;;   :guarded GUARDS       →  :if C1 E1 (:if C2 E2 … (:app error …))
;;   :list-comp EXPR QUALS →  concatMap-based expression (§3.11)
;;
;; Everything else (:app, :op, :lambda, :let, :case, :do, :tuple,
;; :list, :range, :if, :neg, :sect-left / :sect-right, plus all
;; leaf forms and pattern / type nodes) is passed through after
;; recursing into children.

(define
  hk-guards-to-if
  (fn
    (guards)
    (cond
      ((empty? guards)
        (list
          :app
          (list :var "error")
          (list :string "Non-exhaustive guards")))
      (:else
        (let
          ((g (first guards)))
          (list
            :if
            (hk-desugar (nth g 1))
            (hk-desugar (nth g 2))
            (hk-guards-to-if (rest guards))))))))

;; do-notation desugaring (Haskell 98 §3.14):
;;   do { e }              =  e
;;   do { e ; ss }         =  e >> do { ss }
;;   do { p <- e ; ss }    =  e >>= \p -> do { ss }
;;   do { let decls ; ss } =  let decls in do { ss }
(define
  hk-desugar-do
  (fn
    (stmts)
    (cond
      ((empty? stmts) (raise "empty do block"))
      ((empty? (rest stmts))
        (let ((s (first stmts)))
          (cond
            ((= (first s) "do-expr") (hk-desugar (nth s 1)))
            (:else
              (raise "do block must end with an expression")))))
      (:else
        (let
          ((s (first stmts)) (rest-stmts (rest stmts)))
          (let
            ((rest-do (hk-desugar-do rest-stmts)))
            (cond
              ((= (first s) "do-expr")
                (list
                  :app
                  (list
                    :app
                    (list :var ">>")
                    (hk-desugar (nth s 1)))
                  rest-do))
              ((= (first s) "do-bind")
                (list
                  :app
                  (list
                    :app
                    (list :var ">>=")
                    (hk-desugar (nth s 2)))
                  (list :lambda (list (nth s 1)) rest-do)))
              ((= (first s) "do-let")
                (list
                  :let
                  (map hk-desugar (nth s 1))
                  rest-do))
              (:else (raise "unknown do-stmt tag")))))))))

;; List-comprehension desugaring (Haskell 98 §3.11):
;;   [e |             ]  =  [e]
;;   [e | b,     Q    ]  =  if b then [e | Q] else []
;;   [e | p <- l, Q    ]  =  concatMap (\p -> [e | Q]) l
;;   [e | let ds, Q    ]  =  let ds in [e | Q]
(define
  hk-lc-desugar
  (fn
    (e quals)
    (cond
      ((empty? quals) (list :list (list e)))
      (:else
        (let
          ((q (first quals)))
          (let
            ((qtag (first q)))
            (cond
              ((= qtag "q-guard")
                (list
                  :if
                  (hk-desugar (nth q 1))
                  (hk-lc-desugar e (rest quals))
                  (list :list (list))))
              ((= qtag "q-gen")
                (list
                  :app
                  (list
                    :app
                    (list :var "concatMap")
                    (list
                      :lambda
                      (list (nth q 1))
                      (hk-lc-desugar e (rest quals))))
                  (hk-desugar (nth q 2))))
              ((= qtag "q-let")
                (list
                  :let
                  (map hk-desugar (nth q 1))
                  (hk-lc-desugar e (rest quals))))
              (:else
                (raise
                  (str
                    "hk-lc-desugar: unknown qualifier tag "
                    qtag))))))))))

(define
  hk-desugar
  (fn
    (node)
    (cond
      ((not (list? node)) node)
      ((empty? node) node)
      (:else
        (let
          ((tag (first node)))
          (cond
            ((= tag "where")
              (list
                :let (map hk-desugar (nth node 2))
                (hk-desugar (nth node 1))))
            ((= tag "guarded") (hk-guards-to-if (nth node 1)))
            ((= tag "list-comp")
              (hk-lc-desugar (hk-desugar (nth node 1)) (nth node 2)))
            ((= tag "app")
              (list
                :app (hk-desugar (nth node 1))
                (hk-desugar (nth node 2))))
            ((= tag "p-rec")
              (let
                ((cname (nth node 1))
                  (field-pats (nth node 2))
                  (field-order (hk-record-field-names cname)))
                (cond
                  ((nil? field-order)
                    (raise (str "p-rec: no record info for " cname)))
                  (:else
                    (list
                      :p-con
                      cname
                      (map
                        (fn
                          (fname)
                          (let
                            ((p (hk-find-rec-pair field-pats fname)))
                            (cond
                              ((nil? p) (list :p-wild))
                              (:else (hk-desugar (nth p 1))))))
                        field-order))))))
            ((= tag "rec-update")
              (list
                :rec-update
                (hk-desugar (nth node 1))
                (map
                  (fn (p) (list (first p) (hk-desugar (nth p 1))))
                  (nth node 2))))
            ((= tag "rec-create")
              (let
                ((cname (nth node 1))
                  (field-pairs (nth node 2))
                  (field-order (hk-record-field-names cname)))
                (cond
                  ((nil? field-order)
                    (raise (str "rec-create: no record info for " cname)))
                  (:else
                    (let
                      ((acc (list :con cname)))
                      (begin
                        (for-each
                          (fn
                            (fname)
                            (let
                              ((pair
                                  (hk-find-rec-pair field-pairs fname)))
                              (cond
                                ((nil? pair)
                                  (raise
                                    (str
                                      "rec-create: missing field "
                                      fname
                                      " for "
                                      cname)))
                                (:else
                                  (set!
                                    acc
                                    (list
                                      :app
                                      acc
                                      (hk-desugar (nth pair 1))))))))
                          field-order)
                        acc))))))
            ((= tag "op")
              (list
                :op (nth node 1)
                (hk-desugar (nth node 2))
                (hk-desugar (nth node 3))))
            ((= tag "neg") (list :neg (hk-desugar (nth node 1))))
            ((= tag "if")
              (list
                :if (hk-desugar (nth node 1))
                (hk-desugar (nth node 2))
                (hk-desugar (nth node 3))))
            ((= tag "tuple") (list :tuple (map hk-desugar (nth node 1))))
            ((= tag "list") (list :list (map hk-desugar (nth node 1))))
            ((= tag "range")
              (list
                :range (hk-desugar (nth node 1))
                (hk-desugar (nth node 2))))
            ((= tag "range-step")
              (list
                :range-step (hk-desugar (nth node 1))
                (hk-desugar (nth node 2))
                (hk-desugar (nth node 3))))
            ((= tag "lambda")
              (list :lambda (nth node 1) (hk-desugar (nth node 2))))
            ((= tag "let")
              (list
                :let (map hk-desugar (nth node 1))
                (hk-desugar (nth node 2))))
            ((= tag "case")
              (list
                :case (hk-desugar (nth node 1))
                (map hk-desugar (nth node 2))))
            ((= tag "alt")
              (list :alt (hk-desugar (nth node 1)) (hk-desugar (nth node 2))))
            ((= tag "do") (hk-desugar-do (nth node 1)))
            ((= tag "sect-left")
              (list :sect-left (nth node 1) (hk-desugar (nth node 2))))
            ((= tag "sect-right")
              (list :sect-right (nth node 1) (hk-desugar (nth node 2))))
            ((= tag "program")
              (list :program (map hk-desugar (hk-expand-records (nth node 1)))))
            ((= tag "module")
              (list
                :module (nth node 1)
                (nth node 2)
                (nth node 3)
                (map hk-desugar (hk-expand-records (nth node 4)))))
            ((= tag "fun-clause")
              (list
                :fun-clause (nth node 1)
                (map hk-desugar (nth node 2))
                (hk-desugar (nth node 3))))
            ((= tag "instance-decl")
              (list
                :instance-decl (nth node 1)
                (nth node 2)
                (map hk-desugar (nth node 3))))
            ((= tag "pat-bind")
              (list :pat-bind (nth node 1) (hk-desugar (nth node 2))))
            ((= tag "bind")
              (list :bind (nth node 1) (hk-desugar (nth node 2))))
            (:else node)))))))

;; Convenience — tokenize + layout + parse + desugar.
(define hk-record-fields (dict))

(define
  hk-register-record-fields!
  (fn (cname fields) (dict-set! hk-record-fields cname fields)))

(define
  hk-record-field-names
  (fn
    (cname)
    (if (has-key? hk-record-fields cname) (get hk-record-fields cname) nil)))

(define
  hk-record-field-index
  (fn
    (cname fname)
    (let
      ((fields (hk-record-field-names cname)))
      (cond
        ((nil? fields) -1)
        (:else
          (let
            ((i 0) (idx -1))
            (begin
              (for-each
                (fn
                  (f)
                  (begin (when (= f fname) (set! idx i)) (set! i (+ i 1))))
                fields)
              idx)))))))

(define
  hk-find-rec-pair
  (fn
    (pairs name)
    (cond
      ((empty? pairs) nil)
      ((= (first (first pairs)) name) (first pairs))
      (:else (hk-find-rec-pair (rest pairs) name)))))

(define
  hk-record-accessors
  (fn
    (cname rec-fields)
    (let
      ((n (len rec-fields)) (i 0) (out (list)))
      (define
        hk-ra-loop
        (fn
          ()
          (when
            (< i n)
            (let
              ((field (nth rec-fields i)))
              (let
                ((fname (first field)) (j 0) (pats (list)))
                (define
                  hk-pat-loop
                  (fn
                    ()
                    (when
                      (< j n)
                      (begin
                        (append!
                          pats
                          (if
                            (= j i)
                            (list "p-var" "__rec_field")
                            (list "p-wild")))
                        (set! j (+ j 1))
                        (hk-pat-loop)))))
                (hk-pat-loop)
                (append!
                  out
                  (list
                    "fun-clause"
                    fname
                    (list (list "p-con" cname pats))
                    (list "var" "__rec_field")))
                (set! i (+ i 1))
                (hk-ra-loop))))))
      (hk-ra-loop)
      out)))

(define
  hk-expand-records
  (fn
    (decls)
    (let
      ((out (list)))
      (for-each
        (fn
          (d)
          (cond
            ((and (list? d) (= (first d) "data"))
              (let
                ((dname (nth d 1))
                  (tvars (nth d 2))
                  (cons-list (nth d 3))
                  (deriving (if (> (len d) 4) (nth d 4) (list)))
                  (new-cons (list))
                  (accessors (list)))
                (begin
                  (for-each
                    (fn
                      (c)
                      (cond
                        ((= (first c) "con-rec")
                          (let
                            ((cname (nth c 1)) (rec-fields (nth c 2)))
                            (begin
                              (hk-register-record-fields!
                                cname
                                (map (fn (f) (first f)) rec-fields))
                              (append!
                                new-cons
                                (list
                                  "con-def"
                                  cname
                                  (map (fn (f) (nth f 1)) rec-fields)))
                              (for-each
                                (fn (a) (append! accessors a))
                                (hk-record-accessors cname rec-fields)))))
                        (:else (append! new-cons c))))
                    cons-list)
                  (append!
                    out
                    (if
                      (empty? deriving)
                      (list "data" dname tvars new-cons)
                      (list "data" dname tvars new-cons deriving)))
                  (for-each (fn (a) (append! out a)) accessors))))
            (:else (append! out d))))
        decls)
      out)))

(define hk-core (fn (src) (hk-desugar (hk-parse-top src))))

(define hk-core-expr (fn (src) (hk-desugar (hk-parse src))))