;;;;;;;;;;;;;;;;;;;;;;
;; Exprs and Types
;;;;;;;;;;;;;;;;;;;;;;

(datatype Ident)
(datatype Expr)
(datatype Type)
(datatype Scheme)
(datatype Ctx)
;; TODO: can't do unit right now
(sort QuantifiedVs (Set Ident))

(function Fresh (Ident i64) Ident)
(function V (String) Ident)

(function Var (Ident) Expr)
(function App (Expr Expr) Expr)
(function Abs (Ident Expr) Expr)
(function Let (Ident Expr Expr) Expr)
(function Num (i64) Expr)
(function True () Expr)
(function False () Expr)
(function MyUnit () Expr)

(function TVar (Ident) Type :cost 3)
(function TArr (Type Type) Type :cost 1)
(function TInt () Type :cost 1)
(function TBool () Type :cost 1)
(function TUnit () Type :cost 1)

(function Forall (QuantifiedVs Type) Scheme)
(function Nil () Ctx)
(function Cons (Ident Scheme Ctx) Ctx)

(relation ftvCtx (Ctx QuantifiedVs))
(relation ftv (Type QuantifiedVs))
(relation ftvScheme (Scheme QuantifiedVs))
(relation has-qs (Ctx Type QuantifiedVs))
(relation has-qs-demand (Ctx Type))

;;;;;;;;;;;;;;;;;;;;;;
;; Expression size
;;;;;;;;;;;;;;;;;;;;;;
(relation expr-size (Expr i64))
(rule ((= e (Num n)))((expr-size e 1)))
(rule ((= e (Var x)))((expr-size e 1)))
;; asserted facts will be cleared so we define them as rules
(rule ((= e (True))) ((expr-size e 1)))
(rule ((= e (False))) ((expr-size e 1)))
(rule ((= e (MyUnit))) ((expr-size e 1)))
(rule ((= e (App e1 e2)) 
       (expr-size e1 s1) 
       (expr-size e2 s2))
      ((expr-size e (+ (+ s1 s2) 1))))
(rule ((= e (Let x e1 e2)) 
       (expr-size e1 s1) 
       (expr-size e2 s2))
      ((expr-size e (+ (+ s1 s2) 1))))
(rule ((= e (Abs x e1)) 
       (expr-size e1 s1))
      ((expr-size e (+ s1 1))))

;;;;;;;;;;;;;;;;;;;;;;
;; Scheme and Context
;;;;;;;;;;;;;;;;;;;;;;

(rule ((= e (TBool)))((ftv e (set-empty))))
(rule ((= e (TUnit)))((ftv e (set-empty))))
(rule ((= e (TInt)))((ftv e (set-empty))))
(rule ((= e (TVar x)))((ftv e (set-insert (set-empty) x))))
(rule ((= e (TArr fr to))
       (ftv fr s1)
       (ftv to s2))
      ((ftv e (set-union s1 s2))))
(rule ((= c (Nil))) ((ftvCtx c (set-empty))))
(rule ((= e (Forall qs t))
       (ftv t fvs))
      ((ftvScheme e (set-diff fvs qs)))) 
(rule ((= c (Cons x s n))
       (ftvCtx n fvs1)
       (ftvScheme s fvs2))
      ((ftvCtx c (set-union fvs1 fvs2))))


;; TODO: rewrite lookup to use native sets
(function lookup (Ctx Ident) Scheme :cost 1000)
(rewrite (lookup (Cons x s tl) x) s)
(rule (
  (= t (lookup (Cons y s tl) x))
  (!= x y)
)(
  (union t (lookup tl x))
))

;;;;;;;;;;;;;;;;;;;;;;
;; Generalization and Instantiation
;;;;;;;;;;;;;;;;;;;;;;

(function generalize (Ctx Type) Scheme :cost 1000)
(function instantiate (Scheme i64) Type :cost 1000)

(rule ((has-qs-demand ctx (TInt)))
      ((has-qs ctx (TInt) (set-empty))))
(rule ((has-qs-demand ctx (TBool)))
      ((has-qs ctx (TBool) (set-empty))))
(rule ((has-qs-demand ctx (TUnit)))
      ((has-qs ctx (TUnit) (set-empty))))

(rule ((has-qs-demand ctx (TArr fr to)))
      ((has-qs-demand ctx fr)
       (has-qs-demand ctx to)))
(rule ((has-qs-demand ctx (TArr fr to))
       (has-qs ctx fr qs1)
       (has-qs ctx to qs2))
      ((has-qs ctx (TArr fr to) (set-union qs1 qs2))))

(rule ((has-qs-demand ctx (TVar x))
       (ftvCtx ctx key-set)
       (set-contains key-set x))
      ((has-qs ctx (TVar x) (set-empty))))
(rule ((has-qs-demand ctx (TVar x))
       (ftvCtx ctx key-set)
       (set-not-contains key-set x))
      ((has-qs ctx (TVar x) (set-insert (set-empty) x))))

(rule ((= sc (generalize ctx t)))
      ((has-qs-demand ctx t)))
(rewrite (generalize ctx t)
         (Forall qs t)
         :when ((has-qs ctx t qs)))

(function subst-fresh (QuantifiedVs Type i64) Type :cost 1000)
(rewrite (subst-fresh vs (TInt) c) (TInt))
(rewrite (subst-fresh vs (TBool) c) (TBool))
(rewrite (subst-fresh vs (TUnit) c) (TUnit))
(rewrite (subst-fresh vs (TArr fr to) c) 
         (TArr (subst-fresh vs fr c) (subst-fresh vs to c)))
(rule ((= otype (subst-fresh vs (TVar s) c))
       (set-contains vs s))
      ((union otype (TVar (Fresh s c)))))
(rule ((= otype (subst-fresh vs (TVar s) c))
       (set-not-contains vs s))
      ((union otype (TVar s))))

(rewrite (instantiate (Forall vs t) c)
         (subst-fresh vs t c))

;;;;;;;;;;;;;;;;;;;;;;
;; Injectivity
;;;;;;;;;;;;;;;;;;;;;;

(rule ((= (TArr fr1 to1) (TArr fr2 to2)))
      ((union fr1 fr2) 
       (union to1 to2)))

;;;;;;;;;;;;;;;;;;;;;;
;; Type inference
;;;;;;;;;;;;;;;;;;;;;;

; ctx |- expr : type
(function typeof (Ctx Expr i64) Type :cost 1000)

;; Basic types (TInt, TBool, TUnit)
(rewrite (typeof ctx (Num x) c) (TInt))
(rewrite (typeof ctx (True) c) (TBool))
(rewrite (typeof ctx (False) c) (TBool))
(rewrite (typeof ctx (MyUnit) c) (TUnit))

; sc = lookup(ctx, x)
; t = instantiate(sc)
; ---------------
; ctx |- x : t
(rewrite (typeof ctx (Var x)  c)
         (instantiate (lookup ctx x) c))

(rewrite 
  (typeof ctx (Abs x e) c)
  (TArr (TVar (Fresh x c)) 
        (typeof (Cons x (Forall (set-empty) (TVar (Fresh x c)))
                      ctx)
                e cc))
    :when ((= cc (+ c 1))))

(rule ((= to (typeof ctx (App e1 e2) c))
       (= c1 (+ c 1))
       (expr-size e1 sz)
       (= c2 (+ c (+ sz 1))))
      ((union (typeof ctx e1 c1) 
              (TArr (typeof ctx e2 c2) to))))

(rewrite (typeof ctx (Let x e1 e2) c)
         (typeof (Cons x (generalize ctx (typeof ctx e1 c1)) 
                       ctx)
                 e2 c2)
    :when ((= c1 (+ c 1))
           (expr-size e1 sz)
           (= c2 (+ c (+ sz 1)))))

;;;;;;;;;;;;;;;;;;;;;;
;; Occurs check
;;;;;;;;;;;;;;;;;;;;;;
(relation occurs-check (Ident Type))
(function errors () Ident)
(rule ((= (TVar x) (TArr fr to)))
      ((occurs-check x fr)
       (occurs-check x to)))
(rule ((occurs-check x (TVar x)))
      ;; ((set (errors) x)))
      ((panic "occurs check fail")))
(rule ((occurs-check x (TArr fr to)))
      ((occurs-check x fr)
       (occurs-check x to)))

(relation base-types (Type))
(base-types (TInt))
(base-types (TBool))
(base-types (TUnit))
(rule ((base-types t)
       (= t (TArr fr to)))
      ((panic "Unifying base types with functions")) )
(rule ((= (TInt) (TBool))) ((panic "Unifying base types")))
(rule ((= (TInt) (TUnit))) ((panic "Unifying base types")))
(rule ((= (TBool) (TUnit))) ((panic "Unifying base types")))
;;;;;;;;;;;;;;;;;;;;;;
;; TEST
;;;;;;;;;;;;;;;;;;;;;;

(push)
(let id (Abs (V "x") (Var (V "x"))))
(let t-id (typeof (Nil) id 0))
(run 100)
(check (= t-id (TArr (TVar (Fresh (V "x") 0)) (TVar (Fresh (V "x") 0)))))
(pop)

(push)
(let let-poly (Let (V "id") (Abs (V "x") (Var (V "x")))
                  (App (App (Var (V "id")) (Var (V "id")))
                       (App (Var (V "id")) (True)))))
(let t-let-poly (typeof (Nil) let-poly 0))
(run 100)
(check (= t-let-poly (TBool)))
(pop)

(push)
(let id-id (App (Abs (V "x") (Var (V "x")))
                   (Abs (V "y") (Var (V "y")))))
(let t-id-id (typeof (Nil) id-id 0))
(run 100)
(check (= t-id-id (TArr (TVar (Fresh (V "y") 3)) (TVar (Fresh (V "y") 3)))))
(pop)


(push)
(let let-true (Let (V "x") (True) (True)))
(let t-let-true (typeof (Nil) let-true 0))
(run 100)
(check (= t-let-true (TBool)))
(pop)

(push)

(let let-var-true (Let (V "x") (True) (Var (V "x"))))
(let t-let-var-true (typeof (Nil) let-var-true 0))
(run 100)
(check (= t-let-var-true (TBool)))
(pop)

(push)

(let abs-id (Abs (V "x") 
                    (Let (V "y") (Abs (V "z") (Var (V "z"))) (Var (V "y")))))
(let t-abs-id (typeof (Nil) abs-id 0))
(run 100)
(let x (Fresh (V "x") 0))
(let z (Fresh (Fresh (V "z") 2) 4))
(check (= t-abs-id (TArr (TVar x) (TArr (TVar z) (TVar z)))))
(pop)

(push)

(let let-env (Let (V "x") (True) 
                (Let (V "f") (Abs (V "a") (Var (V "a")))
                (Let (V "x") (MyUnit)
                (App (Var (V "f")) (Var (V "x")))
                ))))
(let t-let-env (typeof (Nil) let-env 0))
(run 100)
(check (= t-let-env (TUnit)))
(pop)

(push)
(let let-env-2a (Let (V "x") (MyUnit)
                   (Let (V "f") (Abs (V "y") (Var (V "x")))
                   (Let (V "x") (True)
                   (App (Var (V "f")) (Var (V "x")))))))
(let t-let-env-2a (typeof (Nil) let-env-2a 0))
(run 100)
(check (= t-let-env-2a (TUnit)))
(pop)

(push)

(let let-env-2b (App (Abs (V "x")
                             (Let (V "f") (Abs (V "y") (Var (V "x")))
                             (Let (V "x") (True)
                             (App (Var (V "f")) (Var (V "x"))))))
                        (MyUnit)))
(let t-let-env-2b (typeof (Nil) let-env-2b 0))
(run 100)
(check (= t-let-env-2b (TUnit)))
(pop)

(push)

;; ((lambda (x) ((lambda (f) ((lambda (x) (f x)) #t)) (lambda (y) x))) 5)
(let let-env-hard (App (Abs (V "x") 
                               (App (Abs (V "f") 
                                         (App (Abs (V "x") (App (Var (V "f")) (Var (V "x")))) 
                                              (True)))
                                    (Abs (V "y") (Var (V "x")))))
                          (MyUnit)))
(let t-let-env-hard (typeof (Nil) let-env-hard 0))
(run 100)
(check (= t-let-env-hard (TUnit)))
(pop)

(push)

(let let-inst (Let (V "id") (Abs (V "x") (Var (V "x")))
                 (Let (V "iid") (Abs (V "y") (Var (V "id")))
                 (App (Var (V "iid")) 
                      (App (Var (V "id")) (True))))) )
(let t-let-inst (typeof (Nil) let-inst 0))
(run 100)
(check (= t-let-inst  (TArr (TVar (Fresh (Fresh (Fresh (V "x") 1) 5) 7)) (TVar (Fresh (Fresh (Fresh (V "x") 1) 5) 7)))))
(pop)