;; Haskell 98 layout algorithm (§10.3).
;;
;; Consumes the raw token stream produced by hk-tokenize and inserts
;; virtual braces / semicolons (types vlbrace / vrbrace / vsemi) based
;; on indentation. Newline tokens are consumed and stripped.
;;
;;   (hk-layout (hk-tokenize src))  →  tokens-with-virtual-layout

;; ── Pre-pass ──────────────────────────────────────────────────────
;;
;; Walks the raw token list and emits an augmented stream containing
;; two fresh pseudo-tokens:
;;
;;   {:type "layout-open"   :col N :keyword K}
;;     At stream start (K = "<module>") unless the first real token is
;;     `module` or `{`. Also immediately after every `let` / `where` /
;;     `do` / `of` whose following token is NOT `{`. N is the column
;;     of the token that follows.
;;
;;   {:type "layout-indent" :col N}
;;     Before any token whose line is strictly greater than the line
;;     of the previously emitted real token, EXCEPT when that token
;;     is already preceded by a layout-open (Haskell 98 §10.3 note 3).
;;
;; Raw newline tokens are dropped.

(define
  hk-layout-keyword?
  (fn
    (tok)
    (and
      (= (get tok "type") "reserved")
      (or
        (= (get tok "value") "let")
        (= (get tok "value") "where")
        (= (get tok "value") "do")
        (= (get tok "value") "of")))))

(define
  hk-layout-pre
  (fn
    (tokens)
    (let
      ((result (list))
        (n (len tokens))
        (i 0)
        (prev-line -1)
        (first-real-emitted false)
        (suppress-next-indent false))
      (define
        hk-next-real-idx
        (fn
          (start)
          (let
            ((j start))
            (define
              hk-nri-loop
              (fn
                ()
                (when
                  (and
                    (< j n)
                    (= (get (nth tokens j) "type") "newline"))
                  (do (set! j (+ j 1)) (hk-nri-loop)))))
            (hk-nri-loop)
            j)))
      (define
        hk-pre-step
        (fn
          ()
          (when
            (< i n)
            (let
              ((tok (nth tokens i)) (ty (get tok "type")))
              (cond
                ((= ty "newline") (do (set! i (+ i 1)) (hk-pre-step)))
                (:else
                  (do
                    (when
                      (not first-real-emitted)
                      (do
                        (set! first-real-emitted true)
                        (when
                          (not
                            (or
                              (and
                                (= ty "reserved")
                                (= (get tok "value") "module"))
                              (= ty "lbrace")))
                          (do
                            (append!
                              result
                              {:type "layout-open"
                                :col (get tok "col")
                                :keyword "<module>"
                                :line (get tok "line")})
                            (set! suppress-next-indent true)))))
                    (when
                      (and
                        (>= prev-line 0)
                        (> (get tok "line") prev-line)
                        (not suppress-next-indent))
                      (append!
                        result
                        {:type "layout-indent"
                          :col (get tok "col")
                          :line (get tok "line")}))
                    (set! suppress-next-indent false)
                    (set! prev-line (get tok "line"))
                    (append! result tok)
                    (when
                      (hk-layout-keyword? tok)
                      (let
                        ((j (hk-next-real-idx (+ i 1))))
                        (cond
                          ((>= j n)
                            (do
                              (append!
                                result
                                {:type "layout-open"
                                  :col 0
                                  :keyword (get tok "value")
                                  :line (get tok "line")})
                              (set! suppress-next-indent true)))
                          ((= (get (nth tokens j) "type") "lbrace") nil)
                          (:else
                            (do
                              (append!
                                result
                                {:type "layout-open"
                                  :col (get (nth tokens j) "col")
                                  :keyword (get tok "value")
                                  :line (get tok "line")})
                              (set! suppress-next-indent true))))))
                    (set! i (+ i 1))
                    (hk-pre-step))))))))
      (hk-pre-step)
      result)))

;; ── Main pass: L algorithm ────────────────────────────────────────
;;
;; Stack is a list; the head is the top of stack. Each entry is
;; either the keyword :explicit (pushed by an explicit `{`) or a dict
;; {:col N :keyword K} pushed by a layout-open marker.
;;
;; Rules (following Haskell 98 §10.3):
;;
;;   layout-open(n)  vs stack:
;;     empty or explicit top                 → push n; emit {
;;     n > top-col                           → push n; emit {
;;     otherwise                             → emit { }; retry as indent(n)
;;
;;   layout-indent(n) vs stack:
;;     empty or explicit top                 → drop
;;     n == top-col                          → emit ;
;;     n <  top-col                          → emit }; pop; recurse
;;     n >  top-col                          → drop
;;
;;   lbrace                                  → push :explicit; emit {
;;   rbrace                                  → pop if :explicit; emit }
;;   `in` with implicit let on top           → emit }; pop; emit in
;;   any other token                         → emit
;;
;; EOF: emit } for every remaining implicit context.

(define
  hk-layout-L
  (fn
    (pre-toks)
    (let
      ((result (list))
        (stack (list))
        (n (len pre-toks))
        (i 0))
      (define hk-emit (fn (t) (append! result t)))
      (define
        hk-indent-at
        (fn
          (col line)
          (cond
            ((or (empty? stack) (= (first stack) :explicit)) nil)
            (:else
              (let
                ((top-col (get (first stack) "col")))
                (cond
                  ((= col top-col)
                    (hk-emit
                      {:type "vsemi" :value ";" :line line :col col}))
                  ((< col top-col)
                    (do
                      (hk-emit
                        {:type "vrbrace" :value "}" :line line :col col})
                      (set! stack (rest stack))
                      (hk-indent-at col line)))
                  (:else nil)))))))
      (define
        hk-open-at
        (fn
          (col keyword line)
          (cond
            ((and
                (> col 0)
                (or
                  (empty? stack)
                  (= (first stack) :explicit)
                  (> col (get (first stack) "col"))))
              (do
                (hk-emit
                  {:type "vlbrace" :value "{" :line line :col col})
                (set! stack (cons {:col col :keyword keyword} stack))))
            (:else
              (do
                (hk-emit
                  {:type "vlbrace" :value "{" :line line :col col})
                (hk-emit
                  {:type "vrbrace" :value "}" :line line :col col})
                (hk-indent-at col line))))))
      (define
        hk-close-eof
        (fn
          ()
          (when
            (and
              (not (empty? stack))
              (not (= (first stack) :explicit)))
            (do
              (hk-emit {:type "vrbrace" :value "}" :line 0 :col 0})
              (set! stack (rest stack))
              (hk-close-eof)))))
      ;; Peek past further layout-indent / layout-open markers to find
      ;; the next real token's value when its type is `reserved`.
      ;; Returns nil if no such token.
      (define
        hk-peek-next-reserved
        (fn
          (start)
          (let ((j (+ start 1)) (found nil) (done false))
            (define
              hk-pnr-loop
              (fn
                ()
                (when
                  (and (not done) (< j n))
                  (let
                    ((t (nth pre-toks j)) (ty (get t "type")))
                    (cond
                      ((or
                          (= ty "layout-indent")
                          (= ty "layout-open"))
                        (do (set! j (+ j 1)) (hk-pnr-loop)))
                      ((= ty "reserved")
                        (do (set! found (get t "value")) (set! done true)))
                      (:else (set! done true)))))))
            (hk-pnr-loop)
            found)))
      (define
        hk-layout-step
        (fn
          ()
          (when
            (< i n)
            (let
              ((tok (nth pre-toks i)) (ty (get tok "type")))
              (cond
                ((= ty "eof")
                  (do
                    (hk-close-eof)
                    (hk-emit tok)
                    (set! i (+ i 1))
                    (hk-layout-step)))
                ((= ty "layout-open")
                  (do
                    (hk-open-at
                      (get tok "col")
                      (get tok "keyword")
                      (get tok "line"))
                    (set! i (+ i 1))
                    (hk-layout-step)))
                ((= ty "layout-indent")
                  (cond
                    ((= (hk-peek-next-reserved i) "in")
                      (do (set! i (+ i 1)) (hk-layout-step)))
                    (:else
                      (do
                        (hk-indent-at (get tok "col") (get tok "line"))
                        (set! i (+ i 1))
                        (hk-layout-step)))))
                ((= ty "lbrace")
                  (do
                    (set! stack (cons :explicit stack))
                    (hk-emit tok)
                    (set! i (+ i 1))
                    (hk-layout-step)))
                ((= ty "rbrace")
                  (do
                    (when
                      (and
                        (not (empty? stack))
                        (= (first stack) :explicit))
                      (set! stack (rest stack)))
                    (hk-emit tok)
                    (set! i (+ i 1))
                    (hk-layout-step)))
                ((and
                    (= ty "reserved")
                    (= (get tok "value") "in")
                    (not (empty? stack))
                    (not (= (first stack) :explicit))
                    (= (get (first stack) "keyword") "let"))
                  (do
                    (hk-emit
                      {:type "vrbrace"
                        :value "}"
                        :line (get tok "line")
                        :col (get tok "col")})
                    (set! stack (rest stack))
                    (hk-emit tok)
                    (set! i (+ i 1))
                    (hk-layout-step)))
                (:else
                  (do
                    (hk-emit tok)
                    (set! i (+ i 1))
                    (hk-layout-step))))))))
      (hk-layout-step)
      (hk-close-eof)
      result)))

(define hk-layout (fn (tokens) (hk-layout-L (hk-layout-pre tokens))))