;; 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
            ;; Transformations
            ((= 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)))

            ;; Expression nodes
            ((= tag "app")
              (list
                :app
                (hk-desugar (nth node 1))
                (hk-desugar (nth node 2))))
            ((= 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 (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))))

            ;; Top-level
            ((= tag "program")
              (list :program (map hk-desugar (nth node 1))))
            ((= tag "module")
              (list
                :module
                (nth node 1)
                (nth node 2)
                (nth node 3)
                (map hk-desugar (nth node 4))))

            ;; Decls carrying a body
            ((= tag "fun-clause")
              (list
                :fun-clause
                (nth node 1)
                (nth node 2)
                (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))))

            ;; Everything else: leaf literals, vars, cons, patterns,
            ;; types, imports, type-sigs, data / newtype / fixity, …
            (:else node)))))))

;; Convenience — tokenize + layout + parse + desugar.
(define
  hk-core
  (fn (src) (hk-desugar (hk-parse-top src))))

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