#lang racket
; file contains Unicode characters - download rather than cut/paste from browser

;(provide def check-one reduce-one clear-defs parse-expr strong-reduce type-check type-synth equiv?)
(provide (all-defined-out)); for testing

;; proust-pred-forall-db-holes: for-all, with deBruijn indices and holes
;; Prabhakar Ragde (September 2021)

(require threading) ; not in base install, install threading-lib package via File menu in DrRacket, or raco
(require racket/control) ; for escapes

;; Grammar:

;; expr = (λ x => expr)             ; lambda abstraction
;;      | (expr expr)               ; function application
;;      | (let x = expr in expr)    ; let for local definitions
;;      | x                         ; variable (natural number here)
;;      | (expr : expr)             ; expression annotated with type
;;      | (∀ (x : expr) -> expr)    ; dependent function type
;;      | (expr -> expr)            ; special case equiv to (∀ (_ : expr) -> expr)
;;      | Type                      ; the type of types


;; Structures for Expr (abstract syntax tree representing terms).

(struct Lam (namehint body) #:transparent)
(struct App (rator rand) #:transparent)
(struct PLet (namehint ldef body) #:transparent)
(struct Var (index) #:transparent)
(struct Def (name) #:transparent)
(struct Ann (expr type) #:transparent)
(struct Arrow (namehint domain range) #:transparent)
(struct Type () #:transparent)

;; Binders for holes (part of Expr for interaction but not parsing, no user syntax)

(struct Hole (number type body) #:transparent)
(struct Guess (number type val body) #:transparent)

;; A Context is a (Listof (List Symbol Expr))
;; Association list mapping variables (represented by positions in a list) to expressions
;; The symbol is for printing.

;; Global variables

(define defs empty)
(define deftypes empty)

;; Parsing

;; parse-expr : sexp -> Expr

(define (parse-expr s) (parse-expr-ctx empty s))

;; parse-expr-ctx : Context sexp -> Expr
;; context is used only for names, but this lets us parse on the fly as needed

(define (parse-expr-ctx ctx s)
  (match s
    [`(λ ,(? symbol? x) => ,e) (Lam x (parse-expr-ctx (cons (list x #f) ctx) e))]
    [`(let ,(? symbol? x) = ,de in ,be) (PLet x (parse-expr-ctx ctx de) (parse-expr-ctx (cons (list x #f) ctx) be))]
    [`(∀ (,(? symbol? x) : ,t) -> ,e) (Arrow x (parse-expr-ctx ctx t) (parse-expr-ctx (cons (list x #f) ctx) e))]
    [`(∀ (,(? symbol? x) : ,t) ,(? list? a) ... -> ,e) (Arrow x (parse-expr-ctx ctx t) (parse-expr-ctx  (cons (list x #f) ctx) `(∀ ,@a -> ,e)))]
    [`(,t1 -> ,t2) (Arrow '_ (parse-expr-ctx ctx t1) (parse-expr-ctx (cons (list #f #f) ctx) t2))]
    [`(,t1 -> ,t2 -> ,r ...) (Arrow '_ (parse-expr-ctx ctx t1) (parse-expr-ctx (cons (list #f #f) ctx) `(,t2 -> ,@r)))]
    [`(,e : ,t) (Ann (parse-expr-ctx ctx e) (parse-expr-ctx ctx t))]
    ['Type (Type)]
    [`(,e1 ,e2) (App (parse-expr-ctx ctx e1) (parse-expr-ctx ctx e2))]
    [`(,e1 ,e2 ,e3 ,r ...) (parse-expr-ctx ctx `((,e1 ,e2) ,e3 ,@r))]
    ['_ (parse-error "cannot use underscore\n")]
    [(? symbol? x)
      (cond
        [(list-index x ctx) => Var]
        [(assoc x defs) (Def x)]
        [else (display ctx) (parse-error "unbound variable: ~a\n" x)])]
    [else (parse-error "bad syntax: ~a\n" s)]))

(define (list-index x lst) (list-index-helper x lst 0))

(define (list-index-helper x lst n)
  (cond
    [(empty? lst) #f]
    [(equal? x (first (first lst))) n]
    [else (list-index-helper x (rest lst) (add1 n))]))

(define (parse-error . args)
  (apply error 'parse args))

;; Unparsing, that is, pretty-printing.

;; pretty-print-expr : Context Expr -> String

(define (pretty-print-expr ctx e)
  (match e
    [(Lam x b)
      (define name (freshen x ctx))
      (format "(λ ~a => ~a)" name (pretty-print-expr (cons (list name #f) ctx) b))]
    [(PLet x d b)
      (define name (freshen x ctx))
      (format "(let ~a = ~a in ~a)" name (pretty-print-expr ctx d) (pretty-print-expr (cons (list name #f) ctx) b))]
    [(Hole n t b)
      (format "(?~a : ~a => ~a)" n (pretty-print-expr ctx t) (pretty-print-expr (cons (list n #f) ctx) b))]
    [(Guess n t v b)
      (format "(?~a : ~a ≈ ~a => ~a)" n (pretty-print-expr ctx t) (pretty-print-expr ctx v)
              (pretty-print-expr (cons (list n #f) ctx) b))]
    [(App e1 e2) (format "(~a ~a)" (pretty-print-expr ctx e1) (pretty-print-expr ctx e2))]
    [(Var n) 
       (define name (first (list-ref ctx n)))
       (if (symbol? name) (symbol->string name) (number->string name))]
    [(Def x) (symbol->string x)]
    [(Ann e t) (format "(~a : ~a)" (pretty-print-expr ctx e) (pretty-print-expr ctx t))]
    [(Arrow '_ t1 t2)
      (format "(~a -> ~a)" (pretty-print-expr ctx t1) (pretty-print-expr (cons (list #f #f) ctx) t2))] 
    [(Arrow x t1 t2)
      (define name (freshen x ctx))
      (format "(∀ (~a : ~a) -> ~a)" name (pretty-print-expr ctx t1) (pretty-print-expr (cons (list name #f) ctx) t2))]
    [(Type) "Type"]
    [#f "_"]
    ))

(define (freshen x ctx)
  (if (assoc x ctx)
      (freshen (string->symbol (string-append (symbol->string x) "_")) ctx)
      x))

;; pretty-print-context : Context -> String

(define (pretty-print-context ctx) (first (pretty-print-context-helper ctx)))

;; pretty-print-context-helper : Context -> (list String Context)

(define (pretty-print-context-helper ctx)
  (cond
    [(empty? ctx) (list "" empty)]
    [else
       (match-define (list rst fctx) (pretty-print-context-helper (rest ctx)))
       (define firstname (first (first ctx)))
       (define freshname
         (cond
           [(boolean? firstname) "_"]
           [(number? firstname) (format "?~a" firstname)] ;firstname]
           [else (freshen firstname fctx)]))
       (define firsttype (second (first ctx)))
       (define prtype (if firsttype (pretty-print-expr fctx (second (first ctx))) "_"))
       (list (string-append (format "\n~a : ~a" freshname prtype) rst)
             (cons (list freshname firsttype) fctx))]))

;; substitution using de Bruijn indices

;; based on PLFA (Wadler/Kokke) from an idea by McBride (2005)

;; A Var is a natural number (starting at 0)

;; ext : (Var -> Var) -> (Var -> Var)
;; Given vvmap which maps variables in context Γ to variables in context Δ,
;; produces the resulting map from variables in context Γ,B to variables in context Δ,B

(define (ext vvmap)
  (lambda (v) (if (zero? v) 0 (add1 (vvmap (sub1 v))))))

;; rename : (Var -> Var) Expr -> Expr
;; Given vvmap which maps variables in context Γ to variables in context Δ,
;; and expr in context Γ, produces the resulting expression in context Δ
;; (contexts are implicit, they're not needed in the computation)

(define (rename vvmap expr)
  (match expr
    [(Lam x b) (Lam x (rename (ext vvmap) b))]
    [(PLet x d b) (PLet x (rename vvmap d) (rename (ext vvmap) b))]
    [(Hole n t b) (Hole n (rename vvmap t) (rename (ext vvmap) b))]
    [(Guess n t v b) (Guess n (rename vvmap t) (rename vvmap v) (rename (ext vvmap) b))]
    [(App e1 e2) (App (rename vvmap e1) (rename vvmap e2))]
    [(Var n) (Var (vvmap n))]
    [(Def _) expr]
    [(Ann e t) (Ann (rename vvmap e) (rename vvmap t))]
    [(Arrow x t1 t2) (Arrow x (rename vvmap t1) (rename (ext vvmap) t2))]
    [(Type) expr]))

;; exts : (Var -> Expr) -> (Var -> Expr)
;; Given a map from variables to expressions in context Γ,
;; produces the resulting map from variables to expressions in context Γ,B

(define (exts vtmap)
  (lambda (v) (if (zero? v) (Var 0) (rename add1 (vtmap (sub1 v))))))

;; psubst: (Var -> Expr) Expr -> Expr
;; Given vtmap which maps variables to expressions in context Γ,
;; and expr in context Γ, produces the resulting expression (parallel substitution)

(define (psubst vtmap expr)
  (match expr
    [(Lam x b) (Lam x (psubst (exts vtmap) b))]
    [(PLet x d b) (PLet x (psubst vtmap d) (psubst (exts vtmap) b))]
    [(Hole n t b) (Hole n (psubst vtmap t) (psubst (exts vtmap) b))]
    [(Guess n t v b) (Guess n (psubst vtmap t) (psubst vtmap v) (psubst (exts vtmap) b))]
    [(App e1 e2) (App (psubst vtmap e1) (psubst vtmap e2))]
    [(Var n) (vtmap n)]
    [(Def _) expr]
    [(Ann e t) (Ann (psubst vtmap e) (psubst vtmap t))]
    [(Arrow x t1 t2) (Arrow x (psubst vtmap t1) (psubst (exts vtmap) t2))]
    [(Type) expr]))

;; substitution for (Var 0) and shifting down of all other variables,
;; which is what's needed for beta-reduction

(define (subst repl expr)
  (psubst (lambda (var) (if (zero? var) repl (Var (sub1 var)))) expr))

#|

;; McBride's original version with single traversal function

(struct Kit (vr tm wk) #:transparent) 

(define (trav k tau expr)
  (match expr
    [(Lam x b) (Lam x (trav k (lift k tau) b))]
    [(App e1 e2) (App (trav k tau e1) (trav k tau e2))]
    [(Var n) ((Kit-tm k) (tau n))]
    [(Def _) expr]
    [(Ann e t) (Ann (trav k tau e) (trav k tau t))]
    [(Arrow x t1 t2) (Arrow x (trav k tau t1) (trav k (lift k tau) t2))]
    [(Type) expr]))

(define (lift k tau)
  (lambda (var) (if (zero? var) ((Kit-vr k) var) ((Kit-wk k) (tau (sub1 var))))))

(define (rename rho) (lambda (expr) (trav (Kit identity Var add1) rho expr)))

(define (psubst sigma expr) (trav (Kit Var identity (rename add1)) sigma expr))

(define (subst repl expr) (psubst (lift (Kit (lambda _ repl) Var Var) identity) expr))


;; Version using arithmetic instead of add1/sub1
;; Based on TaPL chapter 7 and the associated OCaml implementation

(define (tmap vtf c expr)
  (define (walk k t)
    (match t
      [(Lam x b) (Lam x (walk (add1 k) b))]
      [(App e1 e2) (App (walk k e1) (walk k e2))]
      [(Var n) (vtf k n)]
      [(Def _) t]
      [(Ann e t) (Ann (walk k e) (walk k t))]
      [(Arrow x t1 t2) (Arrow x (walk k t1) (walk (add1 k) t2))]
      [(Type) t]))
  (walk c expr))

(define (tshift-above d c expr)
  (tmap (lambda (k x) (if (>= x k) (Var (+ x d)) (Var x))) c expr))

(define (tshift d expr) (tshift-above d 0 expr))

(define (tsubst j s t)
  (tmap (lambda (k x) (if (equal? x (+ j k)) (tshift k s) (Var x))) 0 t))

(define (subst s t)
  (tshift -1 (tsubst 0 (tshift 1 s) t)))

|#


;; Erasing hints (because they interfere with equivalence)

(define (erasehints expr)
  (match expr
    [(Lam _ b) (Lam #f (erasehints b))]
    [(PLet _ d b) (PLet #f d b)]
    [(Hole n t b) (Hole n (erasehints t) (erasehints b))]
    [(Guess n t v b) (Guess n (erasehints t) (erasehints v) (erasehints b))]
    [(App e1 e2) (App (erasehints e1) (erasehints e2))]
    [(Var n) (Var n)]
    [(Def _) expr]
    [(Ann e t) (Ann (erasehints e) (erasehints t))]
    [(Arrow _ t1 t2) (Arrow #f (erasehints t1) (erasehints t2))]
    [(Type) expr]))

;; Alpha equivalence (produces #t/#f)

(define (alpha-equiv? e1 e2) (equal? (erasehints e1) (erasehints e2)))

;; Reducing a term

;; reduce : Expr -> Expr
;; Reduces recursively everywhere, especially in bodies of for-alls and lambdas

(define (reduce expr)
  (match expr
    [(Var _) expr]
    [(Def x) (reduce (second (assoc x defs)))]
    [(Hole _ _ _) expr]
    [(Guess _ _ _ _) expr]
    [(Arrow x a b)
       (Arrow x (reduce a) (reduce b))]
    [(Lam x b) (Lam x (reduce b))]
    [(PLet x d b) (reduce (subst d b))]
    [(App f a)
       (define fr (reduce f))
       (define fa (reduce a))
       (match fr
         [(Lam _ b) (reduce (subst fa b))]
         [else (App fr fa)])]
    [(Ann e t) (reduce e)]
    [(Type) (Type)]))

;; Definitional equality

(define (equiv? e1 e2) (alpha-equiv? (reduce e1) (reduce e2)))

;; Context lookup is more complicated because a type is an expression
;; whose indices are relative to the context it was defined in,
;; but the context grew when it was added and maybe subsequently.
;; So the type has to be renamed on the way back from the recursion.

(define (lookup ctx n)
  (cond
    [(empty? ctx) (error 'lookup "empty context, index ~a\n" n)] ; shouldn't happen
    [(zero? n) (rename add1 (second (first ctx)))]
    [else (rename add1 (lookup (rest ctx) (sub1 n)))]))

;; This is the heart of the verifier.
;; type-check and type-synth are mutually-recursive functions that
;;   check an expression has a type in a context, and
;;   synthesize the type of an expression in a context, respectively.
;; They implement bidirectional type checking.

;; type-check : Context Expr Expr -> true/error
;; Produces true if expr has type type in context ctx (or error if not).

(define (type-check ctx expr type)
  (match expr
    [(Lam x b)
       (type-check ctx type (Type))
       (match (reduce type)
         [(Arrow _ tt tw) (type-check (cons (list x tt) ctx) b tw)]
         [else (cannot-check ctx expr type)])]
    [else (if (equiv? (type-synth ctx expr) type) true (cannot-check ctx expr type))]))

;; cannot-check: Context Expr Type -> void
;; Prints generic error message for type-check
;; Error messages can be improved by replacing uses of this with something more specific

(define (cannot-check ctx e t)
  (error 'type-check "cannot typecheck ~a as ~a in context:\n~a"
         (pretty-print-expr ctx e) (pretty-print-expr ctx t) (pretty-print-context ctx)))

;; type-synth : Context Expr -> Type
;; Produces type of expr in context ctx (or error if can't)

(define (type-synth ctx expr)
  (match expr
    [(Var n) (lookup ctx n)] 
    [(Def x)
       (cond
         [(assoc x deftypes) => second]
         [else (cannot-synth ctx expr)])]
    [(Lam x b) (cannot-synth ctx expr)] ; below: shift down from extended context
    [(Hole n t b) (rename sub1 (type-synth (cons t ctx) b))]
    [(Guess n t v b) (rename sub1 (type-synth (cons t ctx) b))]
    [(PLet x d b)
       (define dt (type-synth ctx d))
       (define bt (type-synth (cons dt ctx) b))
       (subst dt bt)]
    [(App f a) 
       (match (reduce (type-synth ctx f))
         [(Arrow _ tt tw) #:when (type-check ctx a tt) (subst a tw)]
         [else (cannot-synth ctx expr)])]
    [(Ann e t) (type-check ctx t (Type)) (type-check ctx e t) t]
    [(Arrow x tt tw)
       (type-check ctx tt (Type))
       (type-check (cons (list x tt) ctx) tw (Type))
       (Type)]
    [(Type) (Type)]))

;; cannot-synth: Context Expr -> void
;; Prints generic error message for type-synth
;; Again, can improve things by being more specific at use sites

(define (cannot-synth ctx expr)
  (error 'type-synth "cannot infer type of ~a in context:\n~a"
         (pretty-print-expr ctx expr) (pretty-print-context ctx)))

;; Definitions

;; def : symbol Expr -> void
;; Side effect is to mutate global defs, deftypes to include new def if it checks out

(define (def name expr)
  (when (assoc name defs) (tactic-error "~a already defined" name))
  (define e (parse-expr expr))
  (define et (type-synth empty e))
  (match e
    [(Ann ex tp) (set! defs (cons (list name ex) defs))
                 (set! deftypes (cons (list name tp) deftypes))]
    [else (set! defs (cons (list name e) defs))
          (set! deftypes (cons (list name et) deftypes))]))

(define (clear-defs) (set! defs empty) (set! deftypes empty))

(define (check-one expr)
  (printf "~a\n" (pretty-print-expr empty (type-synth empty (parse-expr expr)))))

(define (reduce-one expr)
   (printf "~a\n" (pretty-print-expr empty (reduce (parse-expr expr)))))

;; Interactive proof assistant code
;; Based on McBride's OLEG development calculus (PhD thesis, 1999)
;; with implementation ideas from Brady's 2013 Idris paper

;; Local proof state definition

;; contains current proof expr, focussed goal #
(struct PfSt (expr fg) #:transparent)

;; Global state variables

;; current-proof : (union Pfst #f)
(define current-proof #f)

;; current-type : (union Expr #f)
(define current-type #f)

;; hole-ctr: Nat
;; Counter to uniquely label holes
(define hole-ctr 0)

;; use-hole-ctr: -> Nat
;; Produce hole counter value, increment it (for next hole)
(define (use-hole-ctr) (begin0 hole-ctr (set! hole-ctr (add1 hole-ctr))))

;; new-def : sexp -> void
;; Initialize a new definition with a hole and its type

(define (new-def expr)
  (define t (parse-expr expr))
  (type-check empty t (Type))
  (define hn (use-hole-ctr))
  (set! current-proof (PfSt (Hole hn t (Var 0)) hn)) ; ?x : t . x
  (set! current-type t))

;; print-expr : void -> void
;; Prints current-expr and current-type

(define (print-expr)
  (printf "Expr: ~a\nType: ~a\n" (pretty-print-expr empty (PfSt-expr current-proof)) (pretty-print-expr empty current-type)))

;; focus : Nat -> void

(define (focus n)
  (set! current-proof (struct-copy PfSt current-proof [fg n])))

;; tactic-error : for reporting errors in interactive tactic use

(define (tactic-error . args)
  (apply error 'tactic args)) ; or use (abort (void)) for silent return

;; tactic-r : Ctx Expr Type Nat (Ctx Expr -> Expr) -> Expr
;; Applies tac to focussed goal in expr
;; Mostly structural recursion to find hole, but needs to maintain context

(define (tactic-r ctx expr etype fg all? tac)
 (define (tactic-h ctx expr etype) ; fg, all?, tac constant
  (match expr
    [(Hole n t b)
       (cond
         [all? (tac ctx expr) (Hole n t (tactic-h (cons (list n t) ctx) b t))]
         [(equal? n fg) (tac ctx expr)]
         [else (define nt (tactic-h ctx t (Type)))
               (Hole n nt (tactic-h (cons (list n nt) ctx) b etype))])]
    [(Guess n t v b)
       (cond
         [all? (tac ctx expr) (Guess n t v (tactic-h (cons (list n t) ctx) b t))]
         [(equal? n fg) (tac ctx expr)]
         [else (define nt (tactic-h ctx t (Type)))
               (define nv (tactic-h ctx v t))
               (Guess n nt nv (tactic-h (cons (list n nt) ctx) b etype))])]
    [(Lam x b)
       (define ret (reduce etype))
       (match ret
         [(Arrow _ t w) (Lam x (tactic-h (cons (list x t) ctx) b w))] 
         [else (error "wrong lambda type\n" ctx expr etype)])] ; shouldn't happen
    [(PLet x d b)
       (define td (type-synth ctx d))
       (define nd (tactic-h ctx d td))
       (PLet x nd (tactic-h (cons (list x td) ctx) b etype))]
    [(Arrow x t w)
       (define nt (tactic-h ctx t (Type)))
       (Arrow x nt (tactic-h (cons (list x nt) ctx) w (Type)))]
    [(Ann e t) (Ann (tactic-h ctx e t) (tactic-h ctx t (Type)))]
    [(App f a)
       (define tf (type-synth ctx f))
       (define rtf (reduce tf))
       (match rtf
         [(Arrow n t w) (App (tactic-h ctx f rtf) (tactic-h ctx a t))] ; next line shouldn't happen
         [else (error "app type error ~a : ~a\n" (pretty-print-expr ctx expr) (pretty-print-expr ctx etype))])]
    [t t] ; non-recursive cases
    )) ; end of tactic-h
  (tactic-h ctx expr etype))

;; tactic-s : PfSt (Ctx Expr -> PfSt)
;; applies tactic-r to expr/hole in ps

(define (tactic-s ps tac)
  (match-define (PfSt cexpr fg) ps)
  (struct-copy PfSt ps [expr (tactic-r empty cexpr current-type fg #f tac)]))

;; The next functions produce information from a hole other than a replacement expr
;; so they escape when the hole is reached and deliver info, via continuations
;; Currently they work on current-proof, could take proof state as argument if useful

;; do-in-hole : (Context Hole . Args -> X) -> X
;; does f in hole of current-proof, produces result

(define (do-in-hole f . args)
  (call-with-escape-continuation
    (lambda (k)
      (tactic-s current-proof
        (lambda (ctx hole) (k (apply f ctx hole args)))))))

;; get-hole-type : void -> void
;; Produces type of current hole in current-expr

(define (get-hole-type)
  (do-in-hole (lambda (ctx hole) (Hole-type hole))))

;; print-goal : void -> void
;; prints focussed goal type and context for current-expr

(define (print-goal)
  (unless current-proof (tactic-error "no current proof in progress"))
  (unless (PfSt-fg current-proof) (tactic-error "no goal in focus"))
  (do-in-hole
    (lambda (ctx hole)
       (printf "Goal: ?~a : ~a\nContext: ~a\n"
               (Hole-number hole) (pretty-print-expr ctx (Hole-type hole)) (pretty-print-context ctx)))))

;; print-goals : void -> void
;; print all goals for current proof

(define (print-goals)
  (match-define (PfSt cexpr fg) current-proof)
  (tactic-r empty cexpr current-type fg #t
    (lambda (ctx hole)
      (match hole
        [(Hole n t b) (printf "?~a : ~a\n" n (pretty-print-expr ctx t)) hole]
        [(Guess n t v b) (printf "?~a : ~a ≈ ~a\n" n (pretty-print-expr ctx t) (pretty-print-expr ctx v)) hole]
        [else (tactic-error "not a hole ~a\n" hole)]))) ; shouldn't happen
  (void))

;; For when only the context of the hole in current-expr is needed, not type

(define (do-in-hole-ctx f . args)
  (do-in-hole (lambda (ctx hole) (apply f ctx args))))

;; parse-in-hole : sexp -> Expr
;; parses s in context of focussed goal

(define (parse-in-hole s)
  (do-in-hole-ctx parse-expr-ctx s))

;; synth-in-hole : Expr -> Type
;; synthesizes expr in context of focussed goal

(define (synth-in-hole expr)
  (do-in-hole-ctx type-synth expr))

;; print-in-hole : Expr -> void
;; prints expr in context of focussed goal (for debugging)

(define (print-in-hole expr)
  (do-in-hole-ctx pretty-print-expr expr))

;; primitives from McBride's OLEG calculus,
;; applied to focussed-goal of given expr
;; proof state is threaded through these; they may fail or shift focus

;; internal tactics, signalled with /i, will thread state,
;; so they can be combined


;; claim/i : PfSt Type -> PfSt
;; claims t with ?nh : t . ?goal

(define (claim/i ps t)
  (tactic-s ps
    (lambda (ctx hole)
      (unless (Hole? hole) (tactic-error "current goal is not a hole"))
      (type-check ctx t (Type))
      (define nh (use-hole-ctr))
      (Hole nh t (rename add1 hole)))))

;; claim : sexp -> void
;; user tactic, claims se (creates new hole)

(define (claim se)
  (unless current-proof (tactic-error "no current proof in progress"))
  (define t (parse-in-hole se))
  (set! current-proof (claim/i current-proof t)))

;; fill/i : PfSt (union sexp Expr) Boolean -> PfSt [called "try" in OLEG]
;; makes current hole into guess with initial construction v

(define (fill/i ps v)
  (tactic-s ps
    (lambda (ctx hole)
      (unless (Hole? hole) (tactic-error "current goal is not a hole"))
      (match-define (Hole n t b) hole)
      (type-check ctx v t)
      (Guess n t v b))))

;; fill : sexp -> void
;; user tactic, fills current hole with se

(define (fill se)
  (unless current-proof (tactic-error "no current proof in progress"))
  (define v (parse-in-hole se))
  (set! current-proof (fill/i current-proof v)))
    
;; solve/i : PfSt -> PfSt
;; makes guess into solution term

(define (solve/i ps)
  (tactic-s ps
    (lambda (ctx guess)
      (unless (Guess? guess) (printf "ps : ~a\n" ps) (tactic-error "current goal is not a guess"))
      (unless (pure? (Guess-val guess)) (printf "ps : ~a\n" ps) (tactic-error "current guess val contains a hole or guess"))
      (subst (Guess-val guess) (Guess-body guess))))) ; note no result focus

;; pure? : Expr -> boolean
;; checks whether an expression is pure (no holes or guesses)

(define (pure? expr)
  (match expr
    [(Hole _ _ _) #f]
    [(Guess _ _ _ _) #f]
    [(Lam x b) (pure? b)]
    [(PLet _ d b) (and (pure? d) (pure? b))]
    [(Ann e t) (and (pure? e) (pure? t))]
    [(App f a) (and (pure? f) (pure? a))]
    [t #t] ; non-recursive cases
    ))

;; solve : void -> void
;; user tactic, solves current hole with current guess val

(define (solve)
  (unless current-proof (tactic-error "no current proof in progress"))
  (set! current-proof (solve/i current-proof)))

;; attack/i : PfSt -> PfSt
;; makes focussed-goal into form required for introducing lambda, forall, let
;; (has to be ? x : T . x to keep tactics sound)
;; basically a specialized fill

(define (attack/i ps)
  (define nh (use-hole-ctr))
  (define ns
    (tactic-s ps
      (lambda (ctx hole)
        (match hole
          [(Hole n t b) (Guess n t (Hole nh t (Var 0)) b)]
          [else (tactic-error "cannot attack a guess")]))))
  (struct-copy PfSt ns [fg nh])) ; change focus to new hole

;; attack : void -> void
;; user tactic

(define (attack)
  (unless current-proof (tactic-error "no current proof in progress"))
  (set! current-proof (attack/i current-proof)))

;; next tactic (first of many) threads proof state because it changes it
;; but may fail partway through (so can't be allowed to mutate the global state too early)

;; give/i : PfSt Expr -> PfSt
;; applies fill then solve to expression to fill hole exactly

(define (give/i s e)
  (~> s
      (fill/i e)
      (solve/i)))

;; give : sexp -> Expr
;; user tactic for give/i

(define (give se)
  (unless current-proof (tactic-error "no current proof in progress"))
  (define e (parse-in-hole se))
  (set! current-proof (give/i current-proof e)))

;; make-def : symbol -> void
;; make current expr (must be pure) into definition

(define (make-def name)
  (unless current-proof (tactic-error "no current proof in progress"))
  (define current-expr (PfSt-expr current-proof))
  (unless (pure? current-expr) (tactic-error "current expression contains a hole or guess"))
  (type-check empty current-expr current-type) ; just checking
  (printf "~a is:\n" name)
  (print-expr)
  (set! defs (cons (list name current-expr) defs))
  (set! deftypes (cons (list name current-type) deftypes))
  (set! current-proof #f)
  (set! current-type #f)
  (set! hole-ctr 0)
  (printf "No current proof, use new-def to start one\n"))

;; i-lam/i : PfSt symbol -> PfSt
;; introduces a lambda (given namehint nh) in focussed goal
;; hole must be of form (?x : T . x) [otherwise attack]

(define (i-lam/i s nh)
  (define (tac ctx hole) 
    (match hole
      [(Hole n t (Var 0)) 
         (define rt (reduce t)) 
         (match rt
           [(Arrow y s tp) (Lam nh (Hole n tp (Var 0)))] 
           [else (tactic-error "type of goal is ~a, not a function type" t)])]
      [else (tactic-error "current goal must be attacked before lambda intro")]))
  (tactic-s s tac))

;; i-lams/i : PfSt (union symbol (listof symbol)) -> PfSt
;; internal tactic for introducing multiple lambdas given namehints

(define (i-lams/i s nhl)
  (cond
    [(empty? nhl) (tactic-error "please provide at least one name")]
    [(empty? (rest nhl)) (i-lam/i s (first nhl))]
    [else (~> s
              (i-lam/i (first nhl))
              (i-lams/i (rest nhl)))]))

;; i-lams : (union symbol (listof symbol)) -> void
;; user tactic for introducing lambdas given namehints
;; can use (i-lams 'x) or (i-lams '(x y z))
;; does not attack if needed but that could easily be added

(define (i-lams nhl)
  (unless current-proof (tactic-error "no current proof in progress"))
  (unless (or (symbol? nhl) (list? nhl)) (tactic-error "argument is a list of names"))
  (cond
    [(symbol? nhl) (set! current-proof (i-lam/i current-proof nhl))]
    [(empty? nhl) (tactic-error "please provide at least one name")]
    [else (set! current-proof (i-lams/i current-proof nhl))]))

;; i-forall/i : PfSt symbol Type -> PfSt
;; introduces a forall (given namehint and type) in focussed goal of current proof
;; hole must be of form (?x : Type . x) [otherwise attack]

(define (i-forall/i s nh w)
  (define (tac ctx hole)
    (match hole
      [(Hole n t (Var 0)) 
         (type-check ctx t (Type))
         (Arrow nh w (Hole (Hole-number hole) (Type) (Var 0)))]
      [else (tactic-error "current goal must be attacked before lambda intro")]))
  (tactic-s s tac))

;; i-forall : symbol -> void
;; user tactic for introducing forall given namehint
;; again, does not attack if needed

(define (i-forall nhl se)
  (unless current-proof (tactic-error "no current proof in progress"))
  (define w (parse-in-hole se))
  (set! current-proof (i-forall/i current-proof nhl w)))

;; i-let/i : PfSt symbol Expr Expr -> PfSt
;; introduces a let (given namehint, type, value)
;; hole must be of form (?x : T . x) [otherwise attack]

(define (i-let/i s nh t v)
  (define (tac ctx hole)
    (type-check ctx t (Type))
    (type-check ctx v t)
    (PLet nh (Ann v t) (Hole (Hole-number hole) (rename add1 (Hole-type hole)) (Var 0))))
  (tactic-s s tac))

;; i-let : symbol sexp sexp -> void
;; user tactic for introducing let given namehint, type, value
;; does not attack if needed

(define (i-let nh se ve)
  (unless current-proof (tactic-error "no current proof in progress"))
  (define t (parse-in-hole se))
  (define v (parse-in-hole ve))
  (set! current-proof (i-let/i current-proof nh t v)))

;; i-app : sexp -> void
;; introduces an application of function specified by s
;; uses claim, fill, solve (attack not needed)

(define (i-app se)
  (unless current-proof (tactic-error "no current proof in progress"))
  (define f (parse-in-hole se)) 
  (define rtf (reduce (synth-in-hole f)))
  (match rtf
    [(Arrow n t w)
       (define next-focus hole-ctr)
       (define ns
         (~> current-proof
             (claim/i t) ; hole added here will be focus at end
             (fill/i (App (rename add1 f) (Var 0))) ; doesn't need parse
             (solve/i)))
       (set! current-proof (struct-copy PfSt ns [fg next-focus]))]
    [else (tactic-error "not a function type\n")]))

;; refine/i : PfSt Expr Type Type -> PfSt
;; f is expression to be tried, t its current type, dt desired type
;; If f doesn't work as is, add as many holes as needed

(define (refine/i s f t dt)
  (cond
    [(equiv? t dt) (give/i s f)]
    [else
       (match t
         [(Arrow n d r)
            (~> s
                (claim/i d) ; subtlety below, dt under new binder but r out of old one
                (refine/i (App (rename add1 f) (Var 0)) r (rename add1 dt)))]
         [else (tactic-error "cannot refine")])]))

;; refine : sexp -> void
;; user tactic for refinement

(define (refine se)
  (unless current-proof (tactic-error "no current proof in progress"))
  (define f (parse-in-hole se))
  (define rtf (reduce (synth-in-hole f)))
  (define ht (get-hole-type)) 
  (set! current-proof
     (struct-copy PfSt (refine/i current-proof f rtf ht) [fg #f])))

; uncomment for local testing
;(require racket/trace)
;(trace type-check type-synth)

#;(def 'id-id
  '((let id = ((λ x => (λ y => y)) : (∀ (x : Type) -> (x -> x)))
     in ((id (∀ (x : Type) -> (x -> x))) id))
    : (∀ (x : Type) -> (x -> x))))

;; Examples

#;(def 'id-type '(∀ (A : Type) -> (A -> A)))

;; Making id-type
;; Easier to directly write the above,
;; but this illustrates that interaction is possible to make types

#;(begin
 (new-def 'Type)
 (i-forall 'A 'Type)
 (i-forall '_ 'A)
 (give 'A)
 (make-def 'id-type))

; making id

#;(begin
  (new-def 'id-type)
  (i-lams '(A a))
  (give 'a) 
  (make-def 'id)
  )

; modus ponens

#;(def 'mp-type '(∀ (A : Type) -> (∀ (B : Type) -> (A -> ((A -> B) -> B)))))

; making test2

#;(begin
  (new-def 'mp-type)
  (i-lams '(A B a f))
  (i-app 'f)
  (give 'a)
  (make-def 'mp)
  )

; transitivity of implication

#;(def 'trans-type '(∀ (A : Type) -> (∀ (B : Type) ->( ∀ (C : Type) ->
                     ((B -> C) -> ((A -> B) -> (A -> C)))))))

; making trans

#;(begin
  (new-def 'trans-type)
  (i-lams '(A B C f g x))
  (i-app 'f) 
  (i-app 'g)
  (give 'x)
  (make-def 'trans)
  )

#;(def 'impex1-type '(∀ (A : Type) -> (∀ (B : Type) ->( ∀ (C : Type) ->
                     ((A -> B -> C) -> (A -> B) -> (A -> C))))))

#;(begin
  (new-def 'impex1-type)
  (i-lams '(A B C f g a))
  (i-app '(f a))
  (give '(g a))
  (make-def 'impex1)
  )

;; next example is somewhat awkward to solve with i-app
;; need to use named let to construct subformulas to go in function position
  
(def 'impex2-type '(∀ (A : Type) -> (∀ (B : Type) ->( ∀ (C : Type) ->
                     (((A -> B) -> (A -> C)) -> (A -> B -> C))))))

#;(begin
  (new-def 'impex2-type)
  (i-lams '(A B C f a b))
  (i-let 'g '(A -> B) '(λ x => b)) ; awkward!
  (i-let 'h '(A -> C) '(f g))
  (i-app 'h)
  (give 'a)
  (make-def 'impex2-type))

;; better
#;(begin
  (new-def 'impex2-type)
  (i-lams '(A B C f a b))
; (claim '(A -> B)) ; hole 1
; (claim 'A) ; hole 2
; need to apply f to holes 1 and 2! Instead...
  (refine 'f)
  (print-goals)
; Expr: (λ A => (λ B => (λ C => (λ f => (λ a => (λ b => (?1 : (A -> B) => (?2 : A => ((f 1) 2)))))))))
  (focus 1)
  (attack)  ; hole 1 now guess with hole 3 val, hole 3 in focus
; Expr: (λ A => (λ B => (λ C => (λ f => (λ a => (λ b => (?1 : (A -> B) ≈ (?3 : (A -> B) => 3) => (?2 : A => ((f 1) 2)))))))))
  (i-lams '_) ; puts constant λ in hole 3
  (give 'b) ; completes and removes hole 3
  (focus 1) ; hole 1 guess has been purified, time to solve
  (solve)
; Expr: (λ A => (λ B => (λ C => (λ f => (λ a => (λ b => (?2 : A => ((f (λ _ => b)) 2))))))))
  (focus 2)
  (give 'a)
; Expr: (λ A => (λ B => (λ C => (λ f => (λ a => (λ b => ((f (λ _ => b)) a)))))))
  (make-def 'test4)
  )

; repeated need to refocus after solve/give suggests a hole stack/list in proof state
; can guesses be hidden from user? perhaps always attack, refine/give check for
; pure guess in next focus and automatically solve it (and check again...)