;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;
;; environment.ss
;; Richard Cobbe
;;
;; 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" 1 0)))

  (define-struct env (key-eq? keys vals))
  ;; (Env x y) ::= (make-env (x x -> Boolean) (Listof x) (Listof y))
  ;;   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.)

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

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

  ;; (Env x y) -> 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))]))))))

  ;; (x x -> Boolean) -> (Env x y)
  ;; 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 x) (Listof y) (Env x y) -> (Env x y)
  ;; 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 :: (x x -> Boolean) (Listof x) (Listof y)
  ;;                 -> (Listof x) (Listof y)
  ;; returns copies of keys and vals with duplicates removed; earliest
  ;; 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 x) (Listof y) (Env x y) -> (Env x y)
  ;; 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 x y) (x x -> Boolean) -> (Env x y)
  ;; 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 x y) x (-> z) -> (Union y z)
  ;; lookup key in env; call fk if not found.
  (define lookup
    (opt-lambda (env key
                     [fk
                      (lambda ()
                        (raise (make-exn:env:unbound
                                "unbound identifier"
                                (current-continuation-marks)
                                key)))])
      (let ([key-eq? (env-key-eq? env)])
        (recur loop ([keys (env-keys env)]
                     [vals (env-vals env)])
          (cond
           [(null? keys) (fk)]
           [(key-eq? key (car keys)) (car vals)]
           [else (loop (cdr keys) (cdr vals))])))))

  ;; lookup/eq :: (Env x y) x (x x -> Boolean) (-> z) (Union y z)
  (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)))))

  ;; (y z -> z) z (Env x y) -> z
  ;; folds f over bindings in environment from right to left
  (define env-foldr
    (lambda (f base env)
      (foldr f base (env-vals env))))

  ;; (y z -> z) z (Env x y) -> z
  ;; folds f over bindings in environment from left to right
  (define env-foldl
    (lambda (f base env)
      (foldl f base (env-vals env))))

  ;; (Env x y) x -> 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 x y) -> (Listof x)
  ;; 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))))

  ;; (Env x y) (x -> Boolean) -> (Env x y)
  ;; 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 x y) -> (Listof (List x y))
  ;; 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 x y) -> (Listof (Pair x y))
  ;; 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?)))

  (define binary-pred/c (-> any/c any/c boolean?))

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

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

  ;; extend-env raises an exn:fail if the two lists are of different lengths.
  ;; (This is a contract violation---why is this not exn:fail:contract?)

  (provide/contract [make-empty-env (opt-> () [binary-pred/c]
                                           env?)]
                    [extend-env (->r ([keys list?]
                                      [vals
                                       (lambda (v)
                                         (and (list? v)
                                              (= (length keys)
                                                 (length v))))]
                                      [env env?])
                                     env?)]
                    [extend-unique (->r ([keys list?]
                                         [vals
                                          (lambda (v)
                                            (and (list? v)
                                                 (= (length keys)
                                                    (length v))))]
                                         [env env?])
                                        env?)]
                    [weaken-env (-> env? binary-pred/c env?)]
                    [lookup (opt-> (env? any/c)
                                   [(-> any)]
                                   any)]
                    [lookup/eq (case-> (-> env? any/c binary-pred/c any)
                                       (-> env? any/c binary-pred/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->sexp (-> env? (listof (list/c any/c any/c)))]
                    [env->alist (-> env? (listof pair?))]
                    [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])]))