;; lib/ocaml/infer.sx — Algorithm W type inference for OCaml-on-SX.
;;
;; Consumes lib/guest/hm.sx (algebra) and lib/guest/match.sx (unify) per
;; the Phase 5 sequencing. The kit ships fresh-tv, generalize,
;; instantiate, and substitution composition; this file assembles the
;; lambda / app / let / if rules of Algorithm W against the OCaml AST.
;;
;; Coverage in this slice (atoms + core forms):
;;   :int :float :string :char :bool :unit :var :fun :app :let :if
;;   :op (with builtin signatures for +, -, *, /, mod, comparisons, &&, ||)
;;
;; Out of scope: pattern matching, tuples, lists (need product/list types
;; first), records, modules, ADTs, let-rec.
;;
;; Inference state:
;;   env     — dict: name → scheme
;;   counter — one-element list (mutable cell) used by hm-fresh-tv
;;
;; Returned value: {:subst S :type T}.

(define ocaml-hm-counter (fn () (list 0)))

(define ocaml-hm-empty-subst (fn () {}))

(define ocaml-hm-builtin-env
  (fn ()
    (let ((int-int-int (hm-arrow (hm-int) (hm-arrow (hm-int) (hm-int))))
          (int-int-bool (hm-arrow (hm-int) (hm-arrow (hm-int) (hm-bool))))
          (bool-bool-bool (hm-arrow (hm-bool) (hm-arrow (hm-bool) (hm-bool))))
          (str-str-str (hm-arrow (hm-string) (hm-arrow (hm-string) (hm-string))))
          (any-any-bool
            (let ((a (hm-tv "a")))
              (hm-scheme (list "a")
                (hm-arrow a (hm-arrow a (hm-bool))))))
          (a->a
            (let ((a (hm-tv "a")))
              (hm-scheme (list "a") (hm-arrow a a)))))
      {"+"  (hm-monotype int-int-int)
       "-"  (hm-monotype int-int-int)
       "*"  (hm-monotype int-int-int)
       "/"  (hm-monotype int-int-int)
       "mod" (hm-monotype int-int-int)
       "%"  (hm-monotype int-int-int)
       "**" (hm-monotype int-int-int)
       "<"  (hm-monotype int-int-bool)
       ">"  (hm-monotype int-int-bool)
       "<=" (hm-monotype int-int-bool)
       ">=" (hm-monotype int-int-bool)
       "=" any-any-bool
       "<>" any-any-bool
       "&&" (hm-monotype bool-bool-bool)
       "||" (hm-monotype bool-bool-bool)
       "^"  (hm-monotype str-str-str)
       "not" (hm-monotype (hm-arrow (hm-bool) (hm-bool)))
       "succ" (hm-monotype (hm-arrow (hm-int) (hm-int)))
       "pred" (hm-monotype (hm-arrow (hm-int) (hm-int)))
       "abs"  (hm-monotype (hm-arrow (hm-int) (hm-int)))})))

(define ocaml-infer (fn (expr env counter) nil))

;; Unify two types; raise on failure. The match.sx unify returns nil on
;; failure so we wrap it for clearer errors.
(define ocaml-hm-unify
  (fn (t1 t2 subst)
    (let ((s2 (unify t1 t2 subst)))
      (cond
        ((= s2 nil)
          (error (str "ocaml-infer: cannot unify " t1 " with " t2)))
        (else s2)))))

;; Look up name; instantiate scheme to a fresh monotype.
(define ocaml-infer-var
  (fn (name env counter)
    (cond
      ((has-key? env name)
        (let ((scheme (get env name)))
          (let ((t (hm-instantiate scheme counter)))
            {:subst {} :type t})))
      (else (error (str "ocaml-infer: unbound variable " name))))))

(define ocaml-infer-app
  (fn (fn-expr arg-expr env counter)
    (let ((r1 (ocaml-infer fn-expr env counter)))
      (let ((s1 (get r1 :subst)) (t1 (get r1 :type)))
        (let ((env2 (hm-apply-env s1 env)))
          (let ((r2 (ocaml-infer arg-expr env2 counter)))
            (let ((s2 (get r2 :subst)) (t2 (get r2 :type)))
              (let ((tv (hm-fresh-tv counter)))
                (let ((s3 (ocaml-hm-unify
                            (hm-apply s2 t1)
                            (hm-arrow t2 tv)
                            (hm-compose s2 s1))))
                  {:subst s3 :type (hm-apply s3 tv)})))))))))

(define ocaml-infer-fun
  (fn (params body env counter)
    (cond
      ((= (len params) 0)
        (error "ocaml-infer: fun without params"))
      ((= (len params) 1)
        (let ((tv (hm-fresh-tv counter)))
          (let ((env2 (assoc env (first params) (hm-monotype tv))))
            (let ((r (ocaml-infer body env2 counter)))
              (let ((s (get r :subst)) (t-body (get r :type)))
                {:subst s
                 :type (hm-arrow (hm-apply s tv) t-body)})))))
      (else
        ;; Curry: fun x y -> e ≡ fun x -> fun y -> e
        (let ((tv (hm-fresh-tv counter)))
          (let ((env2 (assoc env (first params) (hm-monotype tv))))
            (let ((r (ocaml-infer-fun (rest params) body env2 counter)))
              (let ((s (get r :subst)) (t-rest (get r :type)))
                {:subst s
                 :type (hm-arrow (hm-apply s tv) t-rest)}))))))))

(define ocaml-infer-let
  (fn (name params rhs body env counter)
    (let ((rhs-expr (cond
                      ((= (len params) 0) rhs)
                      (else (list :fun params rhs)))))
      (let ((r1 (ocaml-infer rhs-expr env counter)))
        (let ((s1 (get r1 :subst)) (t1 (get r1 :type)))
          (let ((env2 (hm-apply-env s1 env)))
            (let ((scheme (hm-generalize t1 env2)))
              (let ((env3 (assoc env2 name scheme)))
                (let ((r2 (ocaml-infer body env3 counter)))
                  (let ((s2 (get r2 :subst)) (t2 (get r2 :type)))
                    {:subst (hm-compose s2 s1) :type t2}))))))))))

(define ocaml-infer-if
  (fn (c-ast t-ast e-ast env counter)
    (let ((rc (ocaml-infer c-ast env counter)))
      (let ((sc (get rc :subst)) (tc (get rc :type)))
        (let ((sc2 (ocaml-hm-unify tc (hm-bool) sc)))
          (let ((env2 (hm-apply-env sc2 env)))
            (let ((rt (ocaml-infer t-ast env2 counter)))
              (let ((st (get rt :subst)) (tt (get rt :type)))
                (let ((env3 (hm-apply-env st env2)))
                  (let ((re (ocaml-infer e-ast env3 counter)))
                    (let ((se (get re :subst)) (te (get re :type)))
                      (let ((sf (ocaml-hm-unify
                                  (hm-apply se tt)
                                  te
                                  (hm-compose se (hm-compose st sc2)))))
                        {:subst sf
                         :type (hm-apply sf te)}))))))))))))

(set! ocaml-infer
  (fn (expr env counter)
    (let ((tag (nth expr 0)))
      (cond
        ((= tag "int")    {:subst {} :type (hm-int)})
        ((= tag "float")  {:subst {} :type (hm-int)})  ;; treat float as int for now
        ((= tag "string") {:subst {} :type (hm-string)})
        ((= tag "char")   {:subst {} :type (hm-string)})
        ((= tag "bool")   {:subst {} :type (hm-bool)})
        ((= tag "unit")   {:subst {} :type (hm-con "Unit" (list))})
        ((= tag "var")    (ocaml-infer-var (nth expr 1) env counter))
        ((= tag "fun")    (ocaml-infer-fun (nth expr 1) (nth expr 2) env counter))
        ((= tag "app")    (ocaml-infer-app (nth expr 1) (nth expr 2) env counter))
        ((= tag "let")    (ocaml-infer-let (nth expr 1) (nth expr 2)
                                           (nth expr 3) (nth expr 4) env counter))
        ((= tag "if")     (ocaml-infer-if (nth expr 1) (nth expr 2)
                                          (nth expr 3) env counter))
        ((= tag "neg")
          (let ((r (ocaml-infer (nth expr 1) env counter)))
            (let ((s (get r :subst)) (t (get r :type)))
              (let ((s2 (ocaml-hm-unify t (hm-int) s)))
                {:subst s2 :type (hm-int)}))))
        ((= tag "not")
          (let ((r (ocaml-infer (nth expr 1) env counter)))
            (let ((s (get r :subst)) (t (get r :type)))
              (let ((s2 (ocaml-hm-unify t (hm-bool) s)))
                {:subst s2 :type (hm-bool)}))))
        ((= tag "op")
          ;; Treat (:op OP L R) as (:app (:app (:var OP) L) R) — same rule.
          (ocaml-infer
            (list :app (list :app (list :var (nth expr 1)) (nth expr 2)) (nth expr 3))
            env counter))
        (else (error (str "ocaml-infer: unsupported tag " tag)))))))

;; Top-level convenience: parse + infer + render the type.
(define ocaml-type-of
  (fn (src)
    (let ((expr (ocaml-parse src))
          (env (ocaml-hm-builtin-env))
          (counter (ocaml-hm-counter)))
      (let ((r (ocaml-infer expr env counter)))
        (ocaml-hm-format-type (hm-apply (get r :subst) (get r :type)))))))

;; Pretty-print a type as an OCaml-style string for testing. Only handles
;; the constructors we use: Int / Bool / String / Unit / -> / type-vars.
(define ocaml-hm-format-type
  (fn (t)
    (cond
      ((is-var? t) (str "'" (var-name t)))
      ((is-ctor? t)
        (let ((head (ctor-head t)) (args (ctor-args t)))
          (cond
            ((= head "->")
              (let ((a (nth args 0)) (b (nth args 1)))
                (str
                  (cond
                    ((and (is-ctor? a) (= (ctor-head a) "->"))
                      (str "(" (ocaml-hm-format-type a) ")"))
                    (else (ocaml-hm-format-type a)))
                  " -> " (ocaml-hm-format-type b))))
            (else head))))
      (else (str t)))))