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

;; proust-imp-with-holes: adds holes and refinement to proust-imp-basic
;; Prabhakar Ragde (April 2020)

;; grammar:

;; expr = (λ x => expr)             ; lambda abstraction
;;      | (expr expr)               ; function application
;;      | x                         ; variable
;;      | (expr : type)             ; annotated expression
;;      | ?                         ; hole for unwritten code (forms goal)

;; type = X
;;      | (type -> type)            ; implication / function

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

(struct Lam (var body) #:transparent)
(struct App (rator rand) #:transparent)
(struct Ann (expr type) #:transparent)
(struct Hole (num) #:transparent) ; arbitrary numbering

;; Structures for Type (abstract syntax tree representing types)

(struct Arrow (domain range) #:transparent)

;; A Context is a (Listof (List Symbol Type))
;; (Association list mapping propositional variables to types.)

;; Parsing
;; All very straightforward, could write a macro to reduce boilerplate

;; parse-expr : sexp -> Expr

(define (parse-expr s)
  (match s
    [`(λ ,(? symbol? x) => ,e) (Lam x (parse-expr e))]
    [`(,e1 ,e2) (App (parse-expr e1) (parse-expr e2))]
    ['? (Hole #f)] ; not numbered yet
    [(? symbol? x) x]
    [`(,e : ,t) (Ann (parse-expr e) (parse-type t))]
    [`(,e1 ,e2 ,e3 ,r ...) (parse-expr `((,e1 ,e2) ,e3 ,@r))]
    [else (error 'parse "bad syntax: ~a" s)]))

;; parse-type : sexp -> Type

(define (parse-type t)
  (match t
    [`(,t1 -> ,t2) (Arrow (parse-type t1) (parse-type t2))]
    [`(,t1 -> ,t2 -> ,r ...) (Arrow (parse-type t1) (parse-type `(,t2 -> ,@r)))]
    [(? symbol? X) X] 
    [else (error 'parse "bad syntax: ~a\n" t)]))

;; Unparsing, that is, pretty-printing

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

(define (pretty-print-expr e)
  (match e
    [(Lam x b) (format "(λ ~a => ~a)" x (pretty-print-expr b))]
    [(App e1 e2) (format "(~a ~a)" (pretty-print-expr e1) (pretty-print-expr e2))]
    [(? symbol? x) (format "~a" x)]
    [(Ann e t) (format "(~a : ~a)" (pretty-print-expr e) (pretty-print-type t))]
    [(Hole n) (format "?~a" n)]))

;; pretty-print-type : Type -> String

(define (pretty-print-type t)
  (match t
    [(Arrow t1 t2) (format "(~a -> ~a)" (pretty-print-type t1) (pretty-print-type t2))]
    [else t]))

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

(define (pretty-print-context ctx)
  (cond
    [(empty? ctx) ""]
    [else (string-append (format "\n~a : ~a" (first (first ctx)) (pretty-print-type (second (first ctx))))
                         (pretty-print-context (rest ctx)))]))


;; 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 Type -> boolean
;; Produces true if expr has type t in context ctx (or error if not)
;; Possible side effect: adds new goal to goal-table if hole encountered when refining

(define (type-check ctx expr type)
  (match expr
    [(Lam x t) 
       (match type
         [(Arrow tt tw) (type-check (cons `(,x ,tt) ctx) t tw)]
         [else (cannot-check ctx expr type)])]
    [(Hole n) (when refining (hash-set! goal-table n (list type ctx)))
              true]
    [else (if (equal? (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 e) (pretty-print-type t) (pretty-print-context ctx)))

;; type-synth : Context Expr -> Type
;; Produces type of expr in context ctx (or error if can't)
;; All failing cases spelled out rather than put into "else" at end

(define (type-synth ctx expr)
  (match expr
    [(Lam _ _) (cannot-synth ctx expr)]
    [(Hole _) (cannot-synth ctx expr)]
    [(Ann e t) (type-check ctx e t) t]
    [(App f a)
       (define tf (type-synth ctx f))
        (match tf
         [(Arrow tt tw) #:when (type-check ctx a tt) tw]
         [else (cannot-synth ctx expr)])]
    [(? symbol? x)
       (cond
         [(assoc x ctx) => second]
         [else (cannot-synth ctx expr)])]))

;; 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 synth type of ~a in context:\n~a"
         (pretty-print-expr expr) (pretty-print-context ctx)))

;; Interactive proof assistant code

;; State variables

;; current-expr : Expr
(define current-expr #f)

;; goal-table: (Hash Nat (list Type Context))
;; Stores required type of goal, and context
(define goal-table (make-hash))

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

;; use-hole-ctr: -> Nat
;; Side-effect: increments hole-ctr
;; Numbering could be more sophisticated but this works
(define (use-hole-ctr) (begin0 hole-ctr (set! hole-ctr (add1 hole-ctr))))

;; refining : Bool
;; Flag to indicate when refinement of a hole is happening
(define refining #f)

;; Starting an interactive session

;; set-task! : sexp -> void
;; Initialize state variables for new task
;; s should parse to a type

(define (set-task! s)
  (define t (parse-type s))
  (set! goal-table (make-hash))
  (set! hole-ctr 1)
  (set! current-expr (Ann (Hole 0) t))
  (hash-set! goal-table 0 (list t empty))
  (printf "Task is now\n")
  (print-task))

;; Print utilities for interaction

(define (print-task) (printf "~a\n" (pretty-print-expr current-expr)))

;; print-goal : [Nat] -> void
;; With no arguments, prints all goals
;; With argument n, prints goal n

(define print-goal
  (case-lambda
    [() (hash-for-each goal-table (lambda (n g) (print-goal n)))]
    [(n) (match-define (list t ctx)
           (hash-ref goal-table n (lambda () (error 'refine "no goal of that number"))))
         (printf "Goal ~a has type ~a\n" n (pretty-print-type t))
         (unless (empty? ctx)
           (printf "in context~a\n" (pretty-print-context ctx)))]))

;; Refining a hole in the current expression

;; refine : Nat sexp -> void
;; Refine goal n with expression s, but make sure it typechecks first
;; If it does, remove goal n and rewrite current expression

(define (refine n s)
  (match-define (list t ctx)
    (hash-ref goal-table n (lambda () (error 'refine "no goal numbered ~a" n))))
  (define e (parse-expr s))
  (type-check ctx e t) ; first time, just check
  (define en (number-new-holes e))
  (set! refining #t)
  (type-check ctx en t) ; second time, add new goals to goal table
  (set! refining #f)
  (hash-remove! goal-table n)
  (set! current-expr (replace-goal-with n en current-expr))
  (define ngoals (hash-count goal-table))
  (printf "Task with ~a is now\n" (format "~a goal~a" ngoals (if (= ngoals 1) "" "s")))
  (print-task)
  (print-goal))

;; replace-goal-with : Nat Expr Expr -> Expr
;; Replaces goal n with repl in e
;; This is straightforward recursive substitution
;; Could be more efficient but this is easy to understand

(define (replace-goal-with n repl e)
  (match e
    [(Lam x b) (Lam x (replace-goal-with n repl b))]
    [(App e1 e2) (App (replace-goal-with n repl e1) (replace-goal-with n repl e2))]
    [(? symbol? x) x]
    [(Ann e t) (Ann (replace-goal-with n repl e) t)]
    [(Hole m) (if (= m n) repl (Hole m))]))

;; number-new-holes : Expr -> Expr
;; Replace unnumbered holes (Hole #f) with numbered ones
;; Straightforward recursive substitution again

(define (number-new-holes e)
  (match e
    [(Lam x b) (Lam x (number-new-holes b))]
    [(App e1 e2) (App (number-new-holes e1) (number-new-holes e2))]
    [(? symbol? x) x]
    [(Ann e t) (Ann (number-new-holes e) t)]
    [(Hole m) (if m (Hole m) (Hole (use-hole-ctr)))]))

;; examples with scripts

;; curr-sexp : Expr -> sexp
;; Another unparser

(define (curr-sexp) (with-input-from-string (pretty-print-expr current-expr) read))

;; Can uncomment any of these examples and run the program
;; (they are commented out with #; which extends to matching parenthesis)
;; or type the scripts into the REPL one line at a time to experiment with the interaction
;; in latter case, would use (print-task) and (print-goal) for information en route

(define example1
 (begin
  (set-task! '(A -> (A -> B) -> B))
  (refine 0 '(λ x => ?))
  (refine 1 '(λ y => ?))
  (refine 2 '(y x))
  (curr-sexp)))
;; ((λ x => (λ y => (y x))) : (A -> ((A -> B) -> B)))

(define example2
 (begin
  (set-task! '(A -> B -> A))
  (refine 0 '(λ x => ?))
  (refine 1 '(λ y => ?))
  (refine 2 'x)
  (curr-sexp)))
;; ((λ x => (λ y => x)) : (A -> (B -> A)))