(module partition mzscheme
  (require (lib "class.ss")
           (lib "boundmap.ss" "syntax")
           (lib "stx.ss" "syntax"))
  (provide new-bound-partition
           partition%
           #;id:same-marks?
           identifier=-choices)
  
  (define (new-bound-partition)
    (new partition% (relation id:same-marks?)))
  
  (define partition%
    (class object%
      (init relation)

      (define related? relation)
      (field (rep=>num (make-hash-table)))
      (field (obj=>rep (make-hash-table 'weak)))
      (field (reps null))
      (field (next-num 0))
      
      (define/public (get-partition obj)
        (rep->partition (obj->rep obj)))

      (define/public (same-partition? A B)
        (= (get-partition A) (get-partition B)))
      
      (define/public (obj->rep obj)
        (hash-table-get obj=>rep obj (lambda () (obj->rep* obj))))

      (define (obj->rep* obj)
        (let loop ([reps reps])
          (cond [(null? reps)
                 (new-rep obj)]
                [(related? obj (car reps))
                 (hash-table-put! obj=>rep obj (car reps))
                 (car reps)]
                [else
                 (loop (cdr reps))])))

      (define/private (new-rep rep)
        (hash-table-put! rep=>num rep next-num)
        (set! next-num (add1 next-num))
        (set! reps (cons rep reps))
        rep)
      
      (define/private (rep->partition rep)
        (hash-table-get rep=>num rep))

      ;; Nearly useless as it stands
      (define/public (dump)
        (hash-table-for-each 
         rep=>num
         (lambda (k v)
           (printf "~s => ~s~n" k v))))
      
      (super-new)
      ))
  
  ;; bound-partition%
  (define bound-partition%
    (class object%
      ;; numbers : bound-identifier-mapping[identifier => number]
      (define numbers (make-bound-identifier-mapping))
      (define next-number 0)
      
      (define/public (representative stx)
        (datum->syntax-object stx 'representative))
      
      (define/public (get-partition stx)
        (let* ([r (representative stx)]
               [n (bound-identifier-mapping-get numbers r (lambda _ #f))])
          (or n
              (begin0 next-number
                      (bound-identifier-mapping-put! numbers r next-number)
                      (set! next-number (add1 next-number))))))
      (super-new)))
  
  ;; Different identifier relations for highlighting.

  (define (lift/rep id=?)
    (lambda (A B)
      (let ([ra (datum->syntax-object A 'representative)]
            [rb (datum->syntax-object B 'representative)])
        (id=? ra rb))))
  
  (define (lift id=?)
    (lambda (A B)
      (and (identifier? A) (identifier? B) (id=? A B))))
  
  ;; id:same-marks? : syntax syntax -> boolean
  (define id:same-marks?
    (lift/rep bound-identifier=?))

  ;; id:X-module=? : identifier identifier -> boolean
  ;; If both module-imported, do they come from the same module?
  ;; If both top-bound, then same source.
  (define (id:source-module=? a b)
    (let ([ba (identifier-binding a)]
          [bb (identifier-binding b)])
      (cond [(or (eq? 'lexical ba) (eq? 'lexical bb))
             (module-identifier=? a b)]
            [(and (not ba) (not bb))
             #t]
            [(or (not ba) (not bb))
             #f]
            [else
             (eq? (car ba) (car bb))])))
  (define (id:nominal-module=? A B)
    (let ([ba (identifier-binding A)]
          [bb (identifier-binding B)])
      (cond [(or (eq? 'lexical ba) (eq? 'lexical bb))
             (module-identifier=? A B)]
            [(or (not ba) (not bb))
             (and (not ba) (not bb))]
            [else (eq? (caddr ba) (caddr bb))])))
  
  (define (symbolic-identifier=? A B)
    (eq? (syntax-e A) (syntax-e B)))

  (define identifier=-choices
    `(("bound-identifier=?"  . ,bound-identifier=?)
      ("same marks" . ,id:same-marks?)
      ("module-identifier=?" . ,module-identifier=?)
      ("module-or-top-identifier=?" . ,module-or-top-identifier=?)
      ("symbolic-identifier=?" . ,symbolic-identifier=?)
      ("same source module" . ,id:source-module=?)
      ("same nominal module" . ,id:nominal-module=?)
      ))
  
  )