hash table pattern

hash-pattern.rkt

#lang racket/base

(provide hash hash*)

(require racket/match
         (only-in racket/base [hash racket:hash])
         (for-syntax racket/base
                     racket/list))

(define undef (gensym))
(define user-def (gensym))

(define (undef? v)
  (eq? undef v))

(define (user-def? v)
  (eq? user-def v))

(define (hash-remove-safe h k)
  (if (hash-has-key? h k)
      (hash-remove h k)
      h))

(define (hash-remove-safe! h k)
  (when (hash-has-key? h k)
    (hash-remove! h k)))

;; get-full-mode :: syntax?
;;                  (listof identifier?)
;;                  (listof identifier?)
;;                  -> (listof syntax?)
(define-for-syntax (get-full-mode mode k-ids ref-ids)
  (cond
    [(eq? mode #t)
     (list #`(λ ()
               (define seen (hash-copy-clear e #:kind 'mutable))
               (define cnt
                 (+ #,@(for/list ([k-id (in-list k-ids)]
                                  [ref-id (in-list ref-ids)])
                         #`(cond
                             [(or (hash-has-key? seen #,k-id)
                                  (user-def? #,ref-id))
                              0]
                             [else
                              (hash-set! seen #,k-id #t)
                              1]))))
               (= (hash-count e) cnt)))]
    [else '()]))

;; do-hash :: syntax? stx-list? (or/c #t #f syntax?) -> syntax?
(define-for-syntax (do-hash stx kvps mode)
  (define kvp-list (syntax->list kvps))
  (define-values (k-exprs v-pats def-exprs def-ids)
    (for/fold ([k-exprs '()]
               [v-pats '()]
               [def-exprs '()]
               [def-ids '()]
               #:result (values (reverse k-exprs)
                                (reverse v-pats)
                                (reverse def-exprs)
                                (reverse def-ids)))
              ([kvp (in-list kvp-list)])
      (syntax-case kvp ()
        [(k-expr v-pat #:default def-expr)
         (values (cons #'k-expr k-exprs)
                 (cons #'v-pat v-pats)
                 (cons #'def-expr def-exprs)
                 (cons #'user-def def-ids))]
        [(k-expr v-pat)
         (values (cons #'k-expr k-exprs)
                 (cons #'v-pat v-pats)
                 (cons #'undef def-exprs)
                 (cons #'undef def-ids))]
        [_ (raise-syntax-error #f "expect a key-value group" stx kvp)])))

  (define k-ids (generate-temporaries k-exprs))
  (define ref-ids (generate-temporaries k-exprs))

  (with-syntax ([(k-id ...) k-ids]
                [(ref-id ...) ref-ids]
                [(k-expr ...) k-exprs]
                [(v-pat ...) v-pats]
                [(def-expr ...) def-exprs]
                [(def-id ...) def-ids])
    #`(? hash?
         ;; we use let explicitly to prevent macro expander from nesting too much
         (app (λ (e)
                ;; initially assign k-ids and ref-ids to #f, so that
                ;; if section 1 short-circuits, there's no need to
                ;; evaluate all k-exprs and hash-refs
                (let ([k-id #f] ...
                      [ref-id #f] ...)
                  (values
                   ;; SECTION 1: predicate
                   (or (undef? (begin
                                 (set! k-id k-expr)
                                 (set! ref-id (hash-ref e k-id def-id))
                                 ref-id))
                       ...)

                   ;; henceforth, we can assume ref-ids are not undef

                   ;; SECTION 2: full mode predicate
                   #,@(get-full-mode mode k-ids ref-ids)

                   ;; SECTION 3: values
                   ref-id ...

                   ;; SECTION 4: rest
                   #,@(cond
                        [(syntax? mode)
                         ;; TODO: we are inlining heavily here, which could
                         ;; blow up the program size quite a bit.
                         ;; Should this instead be a run-time computation?
                         (list #`(λ ()
                                   (cond
                                     [(immutable? e)
                                      #,(for/fold ([stx #'e])
                                                  ([k (in-list k-ids)])
                                          #`(hash-remove-safe #,stx #,k))]
                                     [else
                                      (define e* (hash-copy e))
                                      #,@(for/list ([k (in-list k-ids)])
                                           #`(hash-remove-safe! e* #,k))
                                      e*])))]
                        [else '()]))))

              ;; SECTION 1
              #f
              ;; SECTION 2
              #,@(cond
                   [(eq? mode #t) (list #'(app (λ (p) (p)) #t))]
                   [else '()])
              ;; SECTION 3
              (app (λ (ref-id)
                     (if (user-def? ref-id)
                         def-expr
                         ref-id))
                   v-pat) ...
              ;; SECTION 4
              #,@(cond
                   [(syntax? mode)
                    (list #`(app (λ (p) (p)) #,mode))]
                   [else '()])))))

(define-match-expander hash*
  (λ (stx)
    (syntax-case stx ()
      [(_ kvp ... #:rest rest-pat)
       (eq? (syntax-e #'rest-pat) '_)
       (do-hash stx #'(kvp ...) #f)]
      [(_ kvp ... #:rest rest-pat) (do-hash stx #'(kvp ...) #'rest-pat)]
      [(_ kvp ... #:full) (do-hash stx #'(kvp ...) #t)]
      [(_ kvp ... #:partial) (do-hash stx #'(kvp ...) #f)]
      [(_ kvp ...) (do-hash stx #'(kvp ...) #f)])))

(define-for-syntax (make-pairs stx xs)
  (let loop ([xs (syntax->list xs)] [acc '()])
    (cond
      [(empty? xs) (reverse acc)]
      [(empty? (rest xs))
       (raise-syntax-error #f "key does not have a value" stx)]
      [else (loop (rest (rest xs)) (cons (list (first xs) (second xs)) acc))])))

(define-match-expander hash
  (λ (stx)
    (syntax-case stx ()
      [(_ stuff ... #:rest rest-pat)
       (with-syntax ([(kvp ...) (make-pairs stx #'(stuff ...))])
         #'(hash* kvp ... #:rest rest-pat))]
      [(_ stuff ... #:full)
       (with-syntax ([(kvp ...) (make-pairs stx #'(stuff ...))])
         #'(hash* kvp ... #:full))]
      [(_ stuff ... #:partial)
       (with-syntax ([(kvp ...) (make-pairs stx #'(stuff ...))])
         #'(hash* kvp ... #:partial))]
      [(_ stuff ...)
       (with-syntax ([(kvp ...) (make-pairs stx #'(stuff ...))])
         #'(hash* kvp ... #:full))]))
  (λ (stx)
    (syntax-case stx ()
      [(_ stuff ...) #'(racket:hash stuff ...)])))

(module+ test
  (require rackunit)
  (test-case "non hash"
    (check-equal? (match 1
                    [(hash* [3 x]) x]
                    [_ 'failed])
                  'failed))

  (test-case "missing key"
    (check-equal? (match (hash 1 2 5 4)
                    [(hash* [3 x]) x]
                    [_ 'failed])
                  'failed))

  (test-case "value pattern matching"
    (check-equal? (match (hash 1 2 3 4)
                    [(hash* [1 (? odd? x)]) x]
                    [_ 'failed])
                  'failed)

    (check-equal? (match (hash 1 2 3 4)
                    [(hash* [1 (? even? x)]) x]
                    [_ 'failed])
                  2))

  (test-case "key expression"
    (check-equal? (match (hash 1 2 3 4)
                    [(hash* [(+ 1 2) x]) x]
                    [_ 'failed])
                  4))

  (test-case "duplicate"
    (check-equal? (match (hash 1 2 3 4)
                    [(hash* [1 x] [1 y]) (list x y)]
                    [_ 'failed])
                  (list 2 2))

    (check-equal? (match (hash 1 2 3 4)
                    [(hash* [1 x] [1 y] #:full) (list x y)]
                    [_ 'failed])
                  'failed)

    (check-equal? (match (hash 1 2 3 4)
                    [(hash* [1 x] [1 y] [3 z] #:full) (list x y z)]
                    [_ 'failed])
                  (list 2 2 4)))

  (test-case "partial matching"
    (check-equal? (match (hash 1 2 3 4)
                    [(hash* [1 x]) x]
                    [_ 'failed])
                  2))

  (test-case "partial matching (multiple)"
    (check-equal? (match (hash 1 2 3 4 5 6)
                    [(hash* [1 x] [5 z]) (list x z)]
                    [_ 'failed])
                  (list 2 6)))

  (test-case "full matching"
    (check-equal? (match (hash 1 2 3 4 5 6)
                    [(hash* [1 x] [3 y] [5 z] #:full) (list x y z)]
                    [_ 'failed])
                  (list 2 4 6)))

  (test-case "full matching failure"
    ;; extra keys
    (check-equal? (match (hash 1 2 3 4 5 6)
                    [(hash* [1 x] [5 z] #:full) (list x z)]
                    [_ 'failed])
                  'failed)

    ;; missing keys
    (check-equal? (match (hash 1 2 3 4 5 6)
                    [(hash* [1 x] [4 y] [5 z] #:full) (list x y z)]
                    [_ 'failed])
                  'failed))

  (test-case "rest matching"
    ;; single
    (check-equal? (match (hash 1 2 3 4 5 6)
                    [(hash* [1 x] [5 z] #:rest (? hash? h)) (list x z h)]
                    [_ 'failed])
                  (list 2 6 (hash 3 4)))

    ;; multiple
    (check-equal? (match (hash 1 2 3 4 5 6)
                    [(hash* [5 z] #:rest (? hash? h)) (list z h)]
                    [_ 'failed])
                  (list 6 (hash 1 2 3 4)))

    ;; nested
    (check-equal? (match (hash 1 2 3 4 5 6)
                    [(hash* [5 z] #:rest (hash* [1 x])) (list x z)]
                    [_ 'failed])
                  (list 2 6)))

  (test-case "rest matching failure"
    ;; rest-pat
    (check-equal? (match (hash 1 2 3 4 5 6)
                    [(hash* [1 x] [5 z] #:rest (? number? h)) (list x z h)]
                    [_ 'failed])
                  'failed)

    ;; extra keys
    (check-equal? (match (hash 1 2 3 4 5 6)
                    [(hash* [4 z] #:rest (? hash? h)) (list z h)]
                    [_ 'failed])
                  'failed))

  (test-case "evaluate only once (at least when hash* is the top-level pattern)"
    ;; partial mode
    (let ([var 0])
      (check-equal? (match (hash -5 -6 1 2 3 4)
                      [(hash* [(begin (set! var (+ var 1))
                                      var)
                               x]
                              [(begin (set! var (+ var 2))
                                      var)
                               y])
                       (list x y)]
                      [_ 'failed])
                    (list 2 4))
      (check-equal? var 3))

    ;; full mode
    (let ([var 0])
      (check-equal? (match (hash 1 2 3 4)
                      [(hash* [(begin (set! var (+ var 1))
                                      var)
                               x]
                              [(begin (set! var (+ var 2))
                                      var)
                               y]
                              #:full)
                       (list x y)]
                      [_ 'failed])
                    (list 2 4))
      (check-equal? var 3))

    ;; rest mode
    (let ([var 0])
      (check-equal? (match (hash -5 -6 1 2 3 4)
                      [(hash* [(begin (set! var (+ var 1))
                                      var)
                               x]
                              [(begin (set! var (+ var 2))
                                      var)
                               y]
                              #:rest (? (λ (h) (= 1 (hash-count h))) h))
                       (list x y h)]
                      [_ 'failed])
                    (list 2 4 (hash -5 -6)))
      (check-equal? var 3)))

  (test-case "default value"
    ;; mismatch
    (check-equal? (match (hash 1 2
                               3 4
                               5 6)
                    [(hash* [1 (list x) #:default (list 42)]) (list x)]
                    [_ 'failed])
                  'failed)

    ;; partial
    (check-equal? (match (hash 1 (list 2)
                               3 (list 4)
                               5 (list 6))
                    [(hash* [2 (list x) #:default (list 42)]
                            [5 (list z) #:default (list -42)])
                     (list x z)]
                    [_ 'failed])
                  (list 42 6))

    (check-equal? (match (hash 0 2)
                    [(hash* [1 x #:default 3] [1 y #:default 4])
                     (list x y)]
                    [_ 'failed])
                  (list 3 4))

    ;; full
    (check-equal? (match (hash 1 (list 2)
                               3 (list 4)
                               5 (list 6))
                    [(hash* [2 w #:default (list 42)]
                            [1 x]
                            [3 y]
                            [5 z #:default (list -42)]
                            #:full)
                     (list w x y z)]
                    [_ 'failed])
                  (list (list 42) (list 2) (list 4) (list 6)))

    (check-equal? (match (hash 1 2)
                    [(hash* [1 x #:default 3] [1 y #:default 4] #:full)
                     (list x y)]
                    [_ 'failed])
                  (list 2 2))

    ;; full failure
    (check-equal? (match (hash 1 (list 2)
                               3 (list 4)
                               5 (list 6))
                    ;; 1 is not matched here
                    [(hash* [2 w #:default (list 42)]
                            [3 y]
                            [5 z #:default (list -42)]
                            #:full)
                     (list w y z)]
                    [_ 'failed])
                  'failed)

    ;; rest
    (check-equal? (match (hash 1 (list 2)
                               3 (list 4)
                               5 (list 6))
                    [(hash* [2 w #:default (list 42)]
                            [3 y]
                            [5 z #:default (list -42)]
                            #:rest h)
                     (list w y z h)]
                    [_ 'failed])
                  (list (list 42) (list 4) (list 6) (hash 1 (list 2)))))

  (test-case "key comparator"
    (let ([b-1 (box 1)]
          [b-2 (box 3)])
      (check-equal? (match (hasheq b-1 2 b-2 4)
                      [(hash* [b-1 x] [b-2 y]) (list x y)]
                      [_ 'failed])
                    (list 2 4))

      (check-equal? (match (hasheq b-1 2 b-2 4)
                      [(hash* [(box 1) x] [(box 3) y]) (list x y)]
                      [_ 'failed])
                    'failed)))

  (test-case "mutability/weakness"
    (check-equal? (match (make-immutable-hash (list (cons 1 2) (cons 3 4)))
                    [(hash* [1 x] #:rest h) (list x h (immutable? h) (hash-strong? h))]
                    [_ 'failed])
                  (list 2 (make-immutable-hash (list (cons 3 4))) #t #t))

    (check-equal? (match (make-hash (list (cons 1 2) (cons 3 4)))
                    [(hash* [1 x] #:rest h) (list x h (immutable? h) (hash-strong? h))]
                    [_ 'failed])
                  (list 2 (make-hash (list (cons 3 4))) #f #t))

    (check-equal? (match (make-weak-hash (list (cons 1 2) (cons 3 4)))
                    [(hash* [1 x] #:rest h) (list x h (immutable? h) (hash-strong? h))]
                    [_ 'failed])
                  (list 2 (make-weak-hash (list (cons 3 4))) #f #f)))

  (test-case "hash"
    (check-equal? (match (hash 1 2 3 4)
                    [(hash 1 x 3 y) (list x y)])
                  (list 2 4))))