(module idmap mzscheme

  (require (lib "contract.ss")
           (lib "boundmap.ss" "syntax")
           (lib "function.ss" "scheme") ;; for curry
           "syntax-errors.scm")

  (define (alistof key/c value/c)
    (listof (cons/c key/c value/c)))

  ;; An IDMap is a (HashTableof Symbol Bind)
  ;; A Bind is (make-bind ID Any)
  (define-struct bind (key value))

  (define idmap? hash-table?)
  (define idset? hash-table?)

  (provide/contract
   [alistof (-> flat-contract/predicate? flat-contract/predicate?
                flat-contract?)]
   [id=? (-> identifier? identifier? boolean?)]

   [idmap? (-> any/c boolean?)]
   [idmap-domain (-> idmap? list?)]
   [empty-idmap (-> idmap?)]
   [idmap-get (opt-> [idmap? identifier?] [(-> any/c)] any/c)]
   [idmap-get! (-> idmap? identifier? (-> any/c) any/c)]
   [alist->idmap (-> (alistof identifier? any/c) idmap?)]
   [idmap->alist (-> idmap? (alistof identifier? any/c))]
   [idmap-member? (-> idmap? identifier? boolean?)]
   [idmap-put-unique! (-> idmap? identifier? any/c void?)]
   [idmap-put! (-> idmap? identifier? any/c void?)]
   [idmap-join (-> (-> any/c any/c any/c) idmap? idmap? idmap?)]
   [idmap-empty? (-> idmap? boolean?)]
   [idmap-for-each (-> idmap? (-> identifier? any/c any) any)]

   [idset? (-> any/c boolean?)]
   [empty-idset (-> idset?)]
   [list->idset (-> (listof identifier?) idset?)]
   [idset-members (-> idset? (listof identifier?))]
   [idset-member? (-> idset? identifier? boolean?)]
   [idset-subset? (-> idset? idset? boolean?)]
   [idset-add-unique! (-> idset? identifier? void?)]
   [idset-add! (-> idset? identifier? void?)]
   [idset-union! (-> idset? idset? void?)]
   [idset-union (-> idset? idset? idset?)])

  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  ;;
  ;;  Primitive functions
  ;;
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

  (define (id=? one two)
    (eq? (syntax-e one) (syntax-e two)))

  (define (idmap-domain table)
    (hash-table-map table (lambda (sym bind) (bind-key bind))))

  (define idmap-get
    (case-lambda
      [(t id) (bind-value (hash-table-get t (syntax-e id)))]
      [(t id f)
       (let/ec return
         (bind-value (hash-table-get t (syntax-e id)
                                     (lambda () (return (f))))))]))

  (define (empty-idmap) (make-hash-table))

  (define (idmap->alist table)
    (hash-table-map
     table
     (lambda (sym bind)
       (cons (bind-key bind) (bind-value bind)))))

  (define (idmap-put! table id value)
    (hash-table-put! table (syntax-e id) (make-bind id value)))

  (define (hash-table-member? table sym)
    (let/ec return
      (hash-table-get table sym (lambda () (return #f)))
      #t))

  (define (idmap-member? table id)
    (hash-table-member? table (syntax-e id)))

  (define (idmap-for-each table f)
    (hash-table-for-each
     table
     (lambda (sym bind)
       (f (bind-key bind) (bind-value bind)))))

  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  ;;
  ;;  Derived functions
  ;;
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

  (define (idmap-empty? table)
    (let/ec return
      (idmap-for-each table (lambda (id v) (return #f)))
      #t))

  (define (idmap-get! idmap id thunk)
    (idmap-get
     idmap id
     (lambda ()
       (let* ([value (thunk)])
         (idmap-put! idmap id value)
         value))))

  (define (alist->idmap alist)
    (let* ([idmap (empty-idmap)])
      (for-each
       (lambda (pair)
         (idmap-put-unique! idmap (car pair) (cdr pair)))
       alist)
      idmap))

  (define (idmap-put-unique! idmap id value)
    (when (idmap-member? idmap id)
      (syntax-error id "duplicate identifier"))
    (idmap-put! idmap id value))

  (define (idmap-join f a b)
    (assert-subset! a b)
    (assert-subset! b a)
    (alist->idmap
     (hash-table-map
      a
      (lambda (id v)
        (cons id (f v (idmap-get b id (lambda () (id-not-found! id)))))))))

  (define (assert-subset! a b)
    (idmap-for-each
     a
     (lambda (id _)
       (unless (idmap-member? b id)
         (id-not-found! id)))))

  (define (id-not-found! id)
    (syntax-error id "id ~s not found where expected" (syntax-e id)))

  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  ;;
  ;;  IDSet functions
  ;;
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

  (define empty-idset empty-idmap)

  (define (list->idset elems)
    (alist->idmap (map (lambda (elem) (cons elem #t)) elems)))

  (define idset-members idmap-domain)

  (define idset-member? idmap-member?)

  (define (idset-subset? one two)
    (andmap (curry idset-member? two) (idset-members one)))

  (define (idset-add-unique! idset elem)
    (idmap-put-unique! idset elem #t))

  (define (idset-add! idset elem)
    (idmap-put! idset elem #t))

  (define (idset-union! one two)
    (idmap-for-each
     two (lambda (id _) (idset-add! one id))))

  (define (idset-union one two)
    (let* ([union (empty-idset)])
      (idset-union! union one)
      (idset-union! union two)))

  )