#lang racket/base

(provide preserve-order
         preserve<?
         sorted-set-elements
         sorted-dict-entries
         sorted-dict-keys-and-values
         sorted-dict-keys)

(require racket/match)
(require (for-syntax racket/base))
(require "record.rkt")
(require "annotation.rkt")
(require racket/set)
(require racket/dict)
(require data/order)
(require (only-in racket/contract any/c))

;;---------------------------------------------------------------------------
;; Order

(define (typecode v)
  (match v
    [(? boolean?) 0]
    [(? single-flonum?) 1]
    [(? double-flonum?) 2]
    [(? integer? x) 3]
    [(? string?) 4]
    [(? bytes?) 5]
    [(? symbol?) 6]
    [(record _ _) 7]
    [(? list?) 8]
    [(? set?) 9]
    [(? dict?) 10]
    [_ 11]))

(define-syntax chain-order
  (syntax-rules ()
    [(_ o) o]
    [(_ o more ...) (match o
                      ['= (chain-order more ...)]
                      [other other])]))

(define (prepare-for-order v)
  (match v
    [(annotated _ _ item) (prepare-for-order item)]
    [_ v]))

(define preserve-order
  (order 'preserve-order
         any/c
         (lambda (a* b*)
           (define a (prepare-for-order a*))
           (define b (prepare-for-order b*))
           (define ta (typecode a))
           (define tb (typecode b))
           (cond [(< ta tb) '<]
                 [(> ta tb) '>]
                 [else (match ta ;; == tb
                         [7 (chain-order
                             (preserve-order (record-label a) (record-label b))
                             (preserve-order (record-fields a)) (preserve-order (record-fields b)))]
                         [8 (match* (a b)
                              [('() '()) '=]
                              [('() _) '<]
                              [(_ '()) '>]
                              [((cons a0 a1) (cons b0 b1))
                               (chain-order (preserve-order a0 b0) (preserve-order a1 b1))])]
                         [9 (preserve-order (sorted-set-elements a) (sorted-set-elements b))]
                         [10 (preserve-order (sorted-dict-keys a) (sorted-dict-keys b))]
                         [_ (datum-order a b)])]))))

(define preserve<? (order-<? preserve-order))

;;---------------------------------------------------------------------------
;; Sorting & cached sorted items

(define set-cache (make-weak-hasheq))
(define dict-entry-cache (make-weak-hasheq))
(define dict-kv-cache (make-weak-hasheq))
(define dict-key-cache (make-weak-hasheq))

(define (sorted-set-elements v)
  (hash-ref! set-cache v (lambda () (sort (set->list v) preserve<?))))

(define (sorted-dict-entries v)
  (hash-ref! dict-entry-cache v (lambda () (sort (dict->list v) preserve<? #:key car))))

(define (sorted-dict-keys-and-values v)
  (hash-ref! dict-kv-cache
             v
             (lambda () (let loop ((xs (sorted-dict-entries v)))
                          (match xs
                            ['() '()]
                            [(cons (cons kk vv) rest) (cons kk (cons vv (loop rest)))])))))

(define (sorted-dict-keys v)
  (hash-ref! dict-kv-cache v (lambda () (map car (sorted-dict-entries v)))))