;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;
;; environment.ss
;; Richard Cobbe
;; Version 3.0
;;
;; Functions that define a standard rib-cage environment data type.
;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(module environment mzscheme

  (require (lib "contract.ss")
           (lib "etc.ss")
           (lib "list.ss")
           (planet "contract-utils.ss" ("cobbe" "contract-utils.plt" 2)))

  (define-struct env (key-eq? keys vals))
  ;; (Env ?a ?b) ::= (make-env (?a ?a -> Boolean*) (Listof ?a) (Listof ?b))
  ;;   side conditions:
  ;;     1) keys should not contain any duplicates under the equality
  ;;        predicate key-eq? .
  ;;     2) keys and vals should be the same length.
  ;;     3) The order of keys and vals is significant.  Not only does a value
  ;;        have to be in the same position in the list as its corresponding
  ;;        key, but keys that were added later should be closer to the head
  ;;        of the list.  (This makes foldl/foldr easier to write.)

  ;; default-unbound-behavior :: (Parameter (?a -> ?b))
  (define default-unbound-behavior
    (make-parameter (lambda (key)
                      (raise (make-exn:env:unbound
                              "unbound identifier"
                              (current-continuation-marks)
                              key)))
                    (lambda (val)
                      (if (and (procedure? val)
                               (procedure-arity-includes? val 1))
                          val
                          (raise (make-exn:fail:contract
                                  "unbound behavior must be unary function"
                                  (current-continuation-marks)))))))

  (define-struct (exn:env:unbound exn:fail:contract) (key) (make-inspector))
  ;; Env-Unbound ::= (make-exn:env:unbound String Continuation-Mark-Set ?a)

  (define-struct (exn:env:shadow exn:fail:contract) (key) #f)
  ;; Env-Shadow ::= (make-exn:env:unbound String Continuation-Mark-Set ?a)

  ;; (Env ?a ?b) -> Boolean
  ;; determines whether side conditions 1 & 2 above hold for env.
  ;; (Can't check condition 3 programmatically.)
  (define valid-env?
    (lambda (env)
      (let ([key-eq? (env-key-eq? env)])
        (and (= (length (env-keys env)) (length (env-vals env)))
             (recur loop ([keys (env-keys env)])
               (cond
                [(null? keys) #t]
                [(memf (lambda (k) (key-eq? (car keys) k)) (cdr keys)) #f]
                [else (loop (cdr keys))]))))))

  ;; (?a ?a -> Boolean*) -> (Env ?a ?b)
  ;; creates an empty environment with the specified equality predicate.
  (define make-empty-env
    (opt-lambda ([key-eq? eq?])
      (make-env key-eq? null null)))

  ;; extend-env :: (Listof ?a) (Listof ?b) (Env ?a ?b) -> (Env ?a ?b)
  ;; adds bindings to env, removing any shadowed bindings
  (define extend-env
    (lambda (keys vals env)
      (let ([key-eq? (env-key-eq? env)]
            [keys (append keys (env-keys env))]
            [vals (append vals (env-vals env))])
        (let-values ([(keys vals)
                      (remove-shadowed key-eq? keys vals)])
          (make-env key-eq? keys vals)))))

  ;; remove-shadowed :: (?a ?a -> Boolean*) (Listof ?a) (Listof ?b)
  ;;                 -> (Listof ?a) (Listof ?b)
  ;; returns copies of keys and vals with duplicates removed; leftmost (i.e.,
  ;; most recently added) copies are preserved
  (define remove-shadowed
    (lambda (key-eq? keys vals)
      (let ([f (lambda (key1)
                 (lambda (key2)
                   (key-eq? key1 key2)))])
        (recur loop ([keys keys]
                     [vals vals]
                     [keys-seen null]
                     [keys-result null]
                     [vals-result null])
          (cond
           [(null? keys)
            (values (reverse keys-result)
                    (reverse vals-result))]
           [(memf (f (car keys)) keys-seen)
            (loop (cdr keys) (cdr vals) keys-seen keys-result vals-result)]
           [else
            (loop (cdr keys) (cdr vals)
                  (cons (car keys) keys-seen)
                  (cons (car keys) keys-result)
                  (cons (car vals) vals-result))])))))

  ;; extend-unique :: (Listof ?a) (Listof ?b) (Env ?a ?b) -> (Env ?a ?b)
  ;; throws exn:env:shadow if there's a duplicate key.
  (define extend-unique
    (lambda (keys0 vals env)
      (let ([key-eq? (env-key-eq? env)])
        (recur loop ([keys keys0])
          (cond
           [(null? keys) (extend-env keys0 vals env)]
           [(bound? env (car keys))
            (raise (make-exn:env:shadow
                    "extend-unique: argument shadows binding in environment"
                    (current-continuation-marks)
                    (car keys)))]
           [(memf (lambda (key) (key-eq? key (car keys))) (cdr keys))
            (raise (make-exn:env:shadow
                    "extend-unique: duplicate key in argument"
                    (current-continuation-marks)
                    (car keys)))]
           [else (loop (cdr keys))])))))

  (define-syntax env-macro
    (syntax-rules ()
      [(_ key-eq? (key val) ...)
       (extend-env (list key ...) (list val ...) (make-empty-env key-eq?))]))

  (define-syntax symbol-env
    (syntax-rules ()
      [(_ (id val) ...)
       (extend-env (list (quote id) ...)
                   (list val ...)
                   (make-empty-env eq?))]))

  ;; weaken-env :: (Env ?a ?b) (?a ?a -> Boolean*) -> (Env ?a ?b)
  ;; weakens env's equality predicate.
  (define weaken-env
    (lambda (env new-eq?)
      (let* ([old-eq? (env-key-eq? env)]
             [wrapped-eq?
              (lambda (x y)
                (let ([old (old-eq? x y)]
                      [new (new-eq? x y)])
                  (if (and old (not new))
                      (raise (make-exn:fail:contract
                             "env's equality predicate improperly weakened"
                             (current-continuation-marks)))
                      new)))])
        (let-values ([(keys vals)
                      (remove-shadowed wrapped-eq?
                                       (env-keys env)
                                       (env-vals env))])
        (make-env wrapped-eq? keys vals)))))

  ;; lookup :: (Env ?a ?b) ?a (?a -> ?c) -> (Union ?b ?c)
  ;; lookup key in env; call fk if not found.
  (define lookup
    (opt-lambda (env key [fk (default-unbound-behavior)])
      (let ([key-eq? (env-key-eq? env)])
        (recur loop ([keys (env-keys env)]
                     [vals (env-vals env)])
          (cond
           [(null? keys) (fk key)]
           [(key-eq? key (car keys)) (car vals)]
           [else (loop (cdr keys) (cdr vals))])))))

  ;; lookup/eq :: (Env ?a ?b) ?a (?a ?a -> Boolean*) (?a -> ?c) (Union ?b ?c)
  (define lookup/eq
    (case-lambda
      [(env key key-eq?) (lookup (weaken-env env key-eq?) key)]
      [(env key key-eq? fk) (lookup (weaken-env env key-eq?) key fk)]))

  ;; env-map :: (?b -> ?c) (Env ?a ?b) -> (Env ?a ?c)
  ;; maps a function over the environment's values
  (define env-map
    (lambda (f env)
      (make-env (env-key-eq? env)
                (env-keys env)
                (map f (env-vals env)))))

  ;; env-foldr :: (?b ?c -> ?c) ?c (Env ?a ?b) -> ?c
  ;; folds f over bindings in environment from right to left
  (define env-foldr
    (lambda (f base env)
      (foldr f base (env-vals env))))

  ;; env-foldl :: (?b ?c -> ?c) ?c (Env ?a ?b) -> ?c
  ;; folds f over bindings in environment from left to right
  (define env-foldl
    (lambda (f base env)
      (foldl f base (env-vals env))))

  ;; bound? :: (Env ?a ?b) ?a -> Boolean
  ;; determines whether key is bound in env
  (define bound?
    (lambda (env key)
      (let ([key-eq? (env-key-eq? env)])
        (and (memf (lambda (k) (key-eq? k key)) (env-keys env))
             #t))))

  ;; env-domain :: (Env ?a ?b) -> (Listof ?a)
  ;; returns list of all keys bound in environment.
  ;; See design-notes.txt for an explanation of why I copy the list (though
  ;; not, of course, the things in it).
  (define env-domain
    (lambda (env)
      (apply list (env-keys env))))

  ;; restrict-domain :: (Env ?a ?b) -> (?a -> Boolean*) -> (Env ?a ?b)
  ;; restricts env's domain to those elements for which pred? is true.
  (define restrict-domain
    (lambda (env pred?)
      (recur loop ([keys (env-keys env)]
                   [vals (env-vals env)]
                   [key-accum null]
                   [val-accum null])
        (cond
         [(null? keys) (make-env (env-key-eq? env)
                                 (reverse key-accum)
                                 (reverse val-accum))]
         [(pred? (car keys)) (loop (cdr keys) (cdr vals)
                                   (cons (car keys) key-accum)
                                   (cons (car vals) val-accum))]
         [else (loop (cdr keys) (cdr vals) key-accum val-accum)]))))

  ;; env->sexp :: (Env ?a ?b) -> (Listof (List ?a ?b))
  ;; converts env to an sexpr for displaying during debugging and testing.
  (define env->sexp
    (lambda (env)
      (map list (env-keys env) (env-vals env))))

  ;; env->alist :: (Env ?a ?b) -> (Listof (Pair ?a ?b))
  ;; converts env to an alist for displaying during debugging and testing.
  (define env->alist
    (lambda (env)
      (map cons (env-keys env) (env-vals env))))

  ;; I'm leaving these in but not exporting them yet until I have a better
  ;; understanding of the pragmatics of contracts.

  (define fast-env/c
    (lambda (key/c val/c)
      (and/c env?
             (struct/c env any/c (listof key/c) (listof val/c)))))

  (define env/c
    (lambda (key/c val/c)
      (and/c (fast-env/c key/c val/c) valid-env?)))

  ;; allow predicates to return any value, in keeping with booleans in Scheme.
  (define binary-pred/c (-> any/c any/c any))

  ;; ======================================================================

  (provide [rename env-macro env]
           symbol-env)

  (define extend/c (->r ([keys list?]
                         [vals
                          (lambda (v)
                            (and (list? v)
                                 (= (length keys)
                                    (length v))))]
                         [env env?])
                        env?))

  (provide/contract [make-empty-env (() [binary-pred/c] . opt-> . env?)]
                    [extend-env extend/c]
                    [extend-unique extend/c]
                    [weaken-env (env? binary-pred/c . -> . env?)]
                    [default-unbound-behavior parameter?]
                    [lookup ((env? any/c) [(any/c . -> . any)] . opt-> . any)]
                    [lookup/eq (case-> (env? any/c binary-pred/c . -> . any)
                                       (env? any/c binary-pred/c
                                             (any/c . -> . any)
                                             . -> .
                                             any))]
                    [env-map ((any/c . -> . any) env? . -> . env?)]
                    [env-foldr ((any/c any/c . -> . any) any/c env?
                                . -> . any)]
                    [env-foldl ((any/c any/c . -> . any) any/c env?
                                . -> . any)]
                    [bound? (env? any/c . -> . boolean?)]
                    [env-domain (env? . -> . list?)]
                    [restrict-domain (env? predicate/c . -> . env?)]
                    [env->alist (env? . -> . (listof pair?))]
                    [env->sexp (env? . -> . (listof (list/c any/c any/c)))]
                    [env? predicate/c]
                    ;;[env/c (-> (union contract? predicate/c)
                    ;;           (union contract? predicate/c)
                    ;;           contract?)]
                    ;;[fast-env/c (-> (union contract? predicate/c)
                    ;;                    (union contract? predicate/c)
                    ;;                contract?)]

                    [struct (exn:env:unbound exn:fail:contract)
                            ([message string?]
                             [continuation-marks continuation-mark-set?]
                             [key any/c])]
                    [struct (exn:env:shadow exn:fail:contract)
                            ([message string?]
                             [continuation-marks continuation-mark-set?]
                             [key any/c])]))