dds/utils.rkt

#lang racket

;;; dds/utils

;;; Various utilities.

(require
 graph
 (for-syntax syntax/parse racket/list))

;;; Typed section.

(module typed typed/racket
  (require typed/graph typed/rackunit typed-compose typed/racket/stream
           (for-syntax syntax/parse racket/list))

  (provide
   Variable VariableMapping GeneralPair

   eval-with eval1-with auto-hash-ref/explicit auto-hash-ref/:
   extract-symbols any->string stringify-variable-mapping string->any
   handle-org-booleans map-sexp read-org-sexp unorg unstringify-pairs
   read-org-variable-mapping unorgv read-symbol-list drop-first-last
   list-sets->list-strings pretty-print-set pretty-print-set-sets
   update-vertices/unweighted update-graph dotit collect-by-key
   collect-by-key/sets ht-values/list->set hash->list/ordered
   multi-split-at lists-transpose in-random)

  (define-type Variable Symbol)
  (define-type (VariableMapping A) (Immutable-HashTable Variable A))

  (: eval-with (-> (VariableMapping Any) Any AnyValues))
  (define (eval-with ht expr)
    (parameterize ([current-namespace (make-base-namespace)])
      (for ([(x val) (in-hash ht)]) (namespace-set-variable-value! x val))
      (eval expr)))

  (: eval1-with (-> (VariableMapping Any) Any Any))
  (define (eval1-with ht expr)
    (call-with-values (λ () (eval-with ht expr))
                      (λ args (car args))))

  (module+ test
    (test-case "eval-with"
      (check-equal? (eval1-with (hash 'a 1 'b 2) '(+ a b 1))
                    4)
      (define ht : (VariableMapping Integer) (hash 'a 1 'b 2))
      (define expr : Any '(+ a b 1))
      (check-equal? (eval1-with ht expr)
                    4)))

  (define-syntax (auto-hash-ref/explicit stx)
    (syntax-parse stx
      [(_ (ht:id xs:id ...) body:expr)
       #`(let #,(for/list ([x (syntax->list #'(xs ...))])
                  #`[#,x (hash-ref ht '#,x)])
           body)]))

  (module+ test
    (test-case "auto-hash-ref/explicit"
      (define mytable #hash((a . 3) (b . 4)))
      (check-equal? (auto-hash-ref/explicit (mytable b a)
                                            (* a b))
                    12)
      (define ht #hash((a . #t) (b . #f)))
      (check-equal? (auto-hash-ref/explicit (ht a b)
                                            (and (not a) b))
                    #f)))

  (define-syntax (auto-hash-ref/: stx)
    (syntax-parse stx
      [(_ ht:id body)
       (let* ([names/: (collect-colons (syntax->datum #'body))])
         #`(let #,(for/list ([x names/:])
                    ;; put x in the same context as body
                    #`[#,(datum->syntax #'body x)
                       (hash-ref ht '#,(strip-colon x))])
             body))]))

  (module+ test
    (test-case "auto-hash-ref/:"
      (define ht1 #hash((x . #t) (y . #t) (t . #f)))
      (define z #t)
      (check-equal? (auto-hash-ref/: ht1
                                     (and :x (not :y) z (or (and :t) :x)))
                    #f)
      (define ht2 #hash((a . 1) (b . 2)))
      (check-equal? (auto-hash-ref/: ht2 (+ :a (* 2 :b)))
                    5)))

  ;;; The helper functions for auto-hash-ref/:.
  (begin-for-syntax
    ;; Collect all the symbols starting with a colon in datum.
    (define (collect-colons datum)
      (remove-duplicates
       (flatten
        (for/list ([token datum])
          (cond
            [(symbol? token)
             (let ([name (symbol->string token)])
               (if (eq? #\: (string-ref name 0))
                   token
                   '()))]
            [(list? token)
             (collect-colons token)]
            [else '()])))))

    ;; Strip the leading colon off x.
    (define (strip-colon x)
      (let ([x-str (symbol->string x)])
        (if (eq? #\: (string-ref x-str 0))
            (string->symbol (substring x-str 1))
            x))))

  (: extract-symbols (-> Any (Listof Symbol)))
  (define (extract-symbols form)
    (: extract-rec (-> Any (Listof Any)))
    (define (extract-rec form)
      (match form
        [(? symbol?) (list form)]
        [(? list?)
         (flatten (for/list : (Listof Any)
                      ([x form])
                    (extract-symbols x)))]
        [else '()]))
    (cast (extract-rec form) (Listof Symbol)))

  (module+ test
    (test-case "extract-symbols"
      (check-equal? (extract-symbols '(1 (2 3) x (y z 3)))
                    '(x y z))))

  (: any->string (-> Any String))
  (define (any->string x)
    (with-output-to-string (λ () (display x))))

  (module+ test
    (test-case "any->string"
      (check-equal? (any->string 'a) "a")
      (check-equal? (any->string '(a 1 (x y))) "(a 1 (x y))")
      (check-equal? (any->string "hello") "hello")))

  (: stringify-variable-mapping (-> (VariableMapping Any) (VariableMapping String)))
  (define (stringify-variable-mapping ht)
    (for/hash : (VariableMapping String)
        ([(key val) (in-hash ht)]) (values key (any->string val))))

  (module+ test
    (test-case "stringify-variable-mapping"
      (define mp (stringify-variable-mapping #hash((a . (and a b)) (b . (not b)))))
      (check-equal? (hash-ref mp 'a) "(and a b)")
      (check-equal? (hash-ref mp 'b) "(not b)")))

  (: string->any (-> String Any))
  (define (string->any str)
    (with-input-from-string str (λ () (read))))

  (module+ test
    (test-case "string->any"
      (check-equal? (string->any "(or b (not a))") '(or b (not a)))
      (check-equal? (string->any "14") 14)))

  ;;; Given a sexp, converts all "#f" to #f and "#t" to #t.
  ;;;
  ;;; When I read Org-mode tables, I pump them through a call to the
  ;;; prin1 because the Elisp sexps seems incompatible with Racket.
  ;;; On the other hand, Racket Booleans seem to upset Elisp a little,
  ;;; so prin1 wraps them in additional double quotes.  This function
  ;;; removes those quotes.
  (: handle-org-booleans (-> Any Any))
  (define/match (handle-org-booleans datum)
    [("#t") #t]
    [("#f") #f]
    [((? list?)) (map handle-org-booleans datum)]
    [ (_) datum])

  (module+ test
    (test-case "handle-org-booleans"
      (check-equal? (handle-org-booleans "#t") #t)
      (check-equal? (handle-org-booleans "#f") #f)
      (check-equal? (handle-org-booleans '("#t" "#f")) '(#t #f))
      (check-equal? (handle-org-booleans "t") "t")))

  (: map-sexp (-> (-> Any Any) Any Any))
  (define (map-sexp func sexp)
    (match sexp
      [(? list?) (map ((curry map-sexp) func) sexp)]
      [datum (func datum)]))

  (module+ test
    (test-case "map-sexp"
      (check-equal? (map-sexp (λ (x) (add1 (cast x Number))) '(1 2 (4 10) 3))
                    '(2 3 (5 11) 4))))

  (: read-org-sexp (-> String Any))
  (define read-org-sexp
    (compose ((curry map-sexp) (match-lambda
                                 [(and (? string?) str) (string->any str)]
                                 [x x]))
             string->any))
  (define unorg read-org-sexp)

  (module+ test
    (test-case "read-org-sexp"
      (check-equal? (read-org-sexp "((\"a\" \"(and a b)\") (\"b\" \"(or b (not a))\"))")
                    '((a (and a b)) (b (or b (not a)))))
      (check-equal? (unorg "(#t \"#t\" \"#t \" '(1 2 \"#f\"))")
                    '(#t #t #t '(1 2 #f)))))

  (define-type (GeneralPair A B) (U (Pair A B) (List A B)))

  (: unstringify-pairs (-> (Listof (GeneralPair String Any))
                           (Listof (GeneralPair Symbol Any))))
  (define (unstringify-pairs pairs)
    (for/list ([p pairs])
      (match p
        [(list key val)
         (cons (string->symbol key) (if (string? val)
                                        (string->any val)
                                        val))]
        [(cons key val)
         (cons (string->symbol key) (if (string? val)
                                        (string->any val)
                                        val))])))

  (module+ test
    (test-case "unstringify-pairs"
      (check-equal? (unstringify-pairs '(("a" . "1") ("b" . "(and a (not b))")))
                    '((a . 1) (b . (and a (not b)))))
      (check-equal? (unstringify-pairs '(("a" . 1) ("b" . "(and a (not b))")))
                    '((a . 1) (b . (and a (not b)))))))

  (: read-org-variable-mapping (-> String (VariableMapping Any)))
  (define read-org-variable-mapping
    (multi-compose
     (λ ([pairs : (Listof (Pair Symbol Any))])
       (make-immutable-hash pairs))
     (λ (sexp)
       (unstringify-pairs (cast sexp (Listof (GeneralPair String Any)))))
     string->any))

  ;;; A synonym for read-org-variable-mapping.
  (define unorgv read-org-variable-mapping)

  (module+ test
    (test-case "read-org-variable-mapping"
      (define m1 (read-org-variable-mapping "((\"a\" \"(and a b)\") (\"b\" \"(or b (not a))\"))"))
      (define m2 (read-org-variable-mapping "((\"a\" . \"(and a b)\") (\"b\" . \"(or b (not a))\"))"))
      (define m3 (unorgv "((\"a\" . \"(and a b)\") (\"b\" . \"(or b (not a))\"))"))
      (check-equal? (hash-ref m1 'a) '(and a b))
      (check-equal? (hash-ref m2 'a) '(and a b))
      (check-equal? (hash-ref m3 'a) '(and a b))
      (check-equal? (hash-ref m1 'b) '(or b (not a)))
      (check-equal? (hash-ref m2 'b) '(or b (not a)))
      (check-equal? (hash-ref m3 'b) '(or b (not a)))))

  (: read-symbol-list (-> String (Listof Symbol)))
  (define (read-symbol-list str)
    (cast (string->any (string-append "(" str ")")) (Listof Symbol)))

  (module+ test
    (test-case "read-symbol-list"
      (check-equal? (read-symbol-list "a b c") '(a b c))))

  (: drop-first-last (-> String String))
  (define (drop-first-last str)
    (substring str 1 (- (string-length str) 1)))

  (module+ test
    (test-case "drop-first-last"
      (check-equal? (drop-first-last "(a b)") "a b")))

  (: list-sets->list-strings (-> (Listof (Setof Any)) (Listof String)))
  (define (list-sets->list-strings lst)
    (map (multi-compose drop-first-last
                        any->string
                        (λ ([x : (Setof Any)])
                          (set->list x))) lst))

  (module+ test
    (test-case "list-sets->list-strings"
      (check-equal? (list-sets->list-strings (list (set 'x 'y) (set 'z) (set) (set 't)))
                    '("x y" "z" "" "t"))))

  (: pretty-print-set (-> (Setof Any) String))
  (define (pretty-print-set s)
    (string-join (sort (set-map s any->string) string<?)))

  (module+ test
    (test-case "pretty-print-set"
      (check-equal? (pretty-print-set (set 'a 'b 1)) "1 a b")))

  (: pretty-print-set-sets (-> (Setof (Setof Any)) String))
  (define (pretty-print-set-sets ms)
    (string-join (for/list ([m ms]) : (Listof String)
                           (format "{~a}" (pretty-print-set m))) ""))

  (module+ test
    (test-case "pretty-print-set-sets"
      (check-equal? (pretty-print-set-sets (set (set 'a 'b) (set 'c))) "{a b}{c}")))

  (define dotit (compose display graphviz))

  (: update-vertices/unweighted (-> Graph (-> Any Any) Graph))
  (define (update-vertices/unweighted gr func)
    (unweighted-graph/directed
     (for/list ([e (in-edges gr)])
       (match-let ([(list u v) e])
         (list (func u) (func v))))))

  (module+ test
    (test-case "update-vertices/unweighted"
      (define gr1 (directed-graph '((a b) (b c))))
      (define gr2 (undirected-graph '((a b) (b c))))
      (define (dbl [x : Any])
        (define x-str (symbol->string (cast x Symbol)))
        (string->symbol (string-append x-str x-str)))
      (define new-gr1 (update-vertices/unweighted gr1 dbl))
      (define new-gr2 (update-vertices/unweighted gr2 dbl))

      (check-false (has-vertex? new-gr1 'a))
      (check-true (has-vertex? new-gr1 'aa))
      (check-false (has-vertex? new-gr1 'b))
      (check-true (has-vertex? new-gr1 'bb))
      (check-false (has-vertex? new-gr1 'c))
      (check-true (has-vertex? new-gr1 'cc))
      (check-true (has-edge? new-gr1 'aa 'bb))
      (check-true (has-edge? new-gr1 'bb 'cc))

      (check-true (has-edge? new-gr2 'aa 'bb))
      (check-true (has-edge? new-gr2 'bb 'aa))
      (check-true (has-edge? new-gr2 'bb 'cc))
      (check-true (has-edge? new-gr2 'cc 'bb))))

  (: update-graph (->* (Graph) (#:v-func (-> Any Any) #:e-func (-> Any Any)) Graph))
  (define (update-graph gr #:v-func [v-func identity] #:e-func [e-func identity])
    (cond
      [(unweighted-graph? gr)
       (unweighted-graph/directed
        (for/list ([e (in-edges gr)]) : (Listof (List Any Any))
                  (match-let ([(list u v) e])
                    (list (v-func u) (v-func v)))))]
      [else
       (weighted-graph/directed
        (for/list ([e (in-edges gr)]) : (Listof (List Any Any Any))
                  (match-let ([(list u v) e])
                    (list (e-func (edge-weight gr u v))
                          (v-func u) (v-func v)))))]))

  (module+ test
    (test-case "update-graph"
      (define gr1 (directed-graph '((a b) (b c))))
      (define gr2 (undirected-graph '((a b) (b c))))
      (define (dbl [x : Any])
        (define x-str (symbol->string (cast x Symbol)))
        (string->symbol (string-append x-str x-str)))
      (define new-gr1-ug (update-graph gr1 #:v-func dbl))
      (define new-gr2-ug (update-graph gr2 #:v-func dbl))
      (define gr3 (weighted-graph/directed '((10 a b) (11 b c))))
      (define new-gr3 (update-graph gr3
                                    #:v-func dbl
                                    #:e-func (λ (x) (* 2 (cast x Number)))))

      (check-false (has-vertex? new-gr1-ug 'a))
      (check-true (has-vertex? new-gr1-ug 'aa))
      (check-false (has-vertex? new-gr1-ug 'b))
      (check-true (has-vertex? new-gr1-ug 'bb))
      (check-false (has-vertex? new-gr1-ug 'c))
      (check-true (has-vertex? new-gr1-ug 'cc))
      (check-true (has-edge? new-gr1-ug 'aa 'bb))
      (check-true (has-edge? new-gr1-ug 'bb 'cc))

      (check-true (has-edge? new-gr2-ug 'aa 'bb))
      (check-true (has-edge? new-gr2-ug 'bb 'aa))
      (check-true (has-edge? new-gr2-ug 'bb 'cc))
      (check-true (has-edge? new-gr2-ug 'cc 'bb))

      (check-true (has-edge? new-gr3 'aa 'bb))
      (check-false (has-edge? new-gr3 'bb 'aa))
      (check-true (has-edge? new-gr3 'bb 'cc))
      (check-false (has-edge? new-gr3 'cc 'bb))
      (check-equal? (edge-weight new-gr3 'aa 'bb) 20)
      (check-equal? (edge-weight new-gr3 'bb 'cc) 22)))

  (: collect-by-key (All (a b) (-> (Listof a) (Listof b)
                                   (Values (Listof a) (Listof (Listof b))))))
  (define (collect-by-key keys vals)
    (for/fold ([ht : (HashTable a (Listof b))
                   (make-immutable-hash)]
               #:result (values (hash-keys ht) (hash-values ht)))
              ([e keys]
               [l vals])
      ((inst hash-update a (Listof b)) ht e (λ (ls) (cons l ls)) (λ () empty))))

  (module+ test
    (test-case "collect-by-key"
      (define-values (e1 l1) (collect-by-key '((1 2) (1 3)) '(a b)))
      (define-values (e2 l2) (collect-by-key '((1 2) (1 2)) '(a b)))
      (check-equal? e1 '((1 2) (1 3))) (check-equal? l1 '((a) (b)))
      (check-equal? e2 '((1 2))) (check-equal? l2 '((b a)))))

  (: collect-by-key/sets (All (a b) (-> (Listof a) (Listof b)
                                        (Values (Listof a) (Listof (Setof b))))))
  (define (collect-by-key/sets edges labels)
    (define-values (es ls) (collect-by-key edges labels))
    (values es ((inst map (Setof b) (Listof b)) list->set ls)))

  (module+ test
    (test-case "collect-by-key/sets"
      (define-values (e3 l3) (collect-by-key/sets '(a b a) '(1 2 1)))
      (check-equal? e3 '(b a)) (check-equal? l3 (list (set 2) (set 1)))))

  ;;; Converts the values of a hash table from lists to sets.
  (: ht-values/list->set (All (a b) (-> (HashTable a (Listof b)) (HashTable a (Setof b)))))
  (define (ht-values/list->set ht)
    (for/hash ([(k v) (in-hash ht)]) : (HashTable a (Setof b))
      (values k (list->set v))))

  (module+ test
    (test-case "ht-values/list->set"
      (check-equal? (ht-values/list->set #hash((a . (1 1))))
                    (hash 'a (set 1)))))

  ;; TODO: Remove after Racket 8.4.
  (: hash->list/ordered (All (a b) (-> (HashTable a b) (Listof (Pairof a b)))))
  (define (hash->list/ordered ht)
    ((inst hash-map a b (Pairof a b)) ht cons #t))

  (module+ test
    (test-case "hash->list/ordered"
      (check-equal? (hash->list/ordered #hash((b . 1) (a . 1)))
                    '((a . 1) (b . 1)))))

  (: multi-split-at (All (a) (-> (Listof (Listof a)) Integer
                                 (Values (Listof (Listof a)) (Listof (Listof a))))))
  (define (multi-split-at lists pos)
    (for/fold ([lefts : (Listof (Listof a)) '()]
               [rights : (Listof (Listof a)) '()]
               #:result (values (reverse lefts) (reverse rights)))
              ([lst (in-list lists)])
      (define-values (left right) ((inst split-at a) lst pos))
      (values (cons left lefts) (cons right rights))))

  (module+ test
    (test-case "multi-split-at"
      (define-values (l1 l2) (multi-split-at '((1 2 3) (a b c)) 2))
      (check-equal? l1 '((1 2) (a b))) (check-equal? l2 '((3) (c)))))

  ;; https://racket.discourse.group/t/get-to-type-apply-in-parallel-lst/683
  (: lists-transpose (All (a ...) (-> (List (Listof a) ... a) (Listof (List a ... a)))))
  (define (lists-transpose lists)
    (sequence->list (in-values-sequence (apply in-parallel lists))))

  (module+ test
  (test-case "lists-transpose"
    (check-equal? (lists-transpose '((1 2) (a b))) '((1 a) (2 b)))))

  (: in-random (case->
              (-> (Sequenceof Flonum))
              (-> Integer (Sequenceof Nonnegative-Fixnum))
              (-> Integer Integer (Sequenceof Nonnegative-Fixnum))))
  (define in-random
    (case-lambda
      [() (stream-cons (random) (in-random))]
      [(k) (stream-cons (random k) (in-random k))]
      [(min max) (stream-cons (random min max) (in-random min max))]))

  (module+ test
    (test-case "in-random"
      (random-seed 1)
      (check-equal? (stream->list (stream-take (in-random 100) 10))
                    '(50 84 10 99 94 88 43 41 63 50))
      (check-equal? (stream->list (stream-take (in-random 50 100) 10))
                    '(57 98 82 83 61 53 73 82 50 80))
      (check-equal? (stream->list (stream-take (in-random) 10))
                    '(0.2718099186980313
                      0.7319496826374751
                      0.17365244033739616
                      0.5593031443038616
                      0.3345256691289459
                      0.9845704615094365
                      0.05753824253751768
                      0.22552976312818723
                      0.21646500425988832
                      0.15188352823997242))))
  )

(require 'typed)
(provide eval-with eval1-with auto-hash-ref/explicit auto-hash-ref/:
         extract-symbols any->string stringify-variable-mapping string->any
         map-sexp read-org-sexp unorg unstringify-pairs
         read-org-variable-mapping unorgv read-symbol-list drop-first-last
         list-sets->list-strings pretty-print-set pretty-print-set-sets
         update-vertices/unweighted update-graph dotit collect-by-key
         collect-by-key/sets ht-values/list->set hash->list/ordered
         multi-split-at lists-transpose in-random)

;;; Untyped section.

(provide
 ;; Functions
 (contract-out [cartesian-product/stream (->* () #:rest (listof stream?) stream?)]
               [boolean-power (-> number? (listof (listof boolean?)))]
               [boolean-power/stream (-> number? (stream/c (listof boolean?)))]
               [any->01 (-> any/c (or/c 0 1))]
               [01->boolean (-> (or/c 0 1) boolean?)])
 ;; Contracts
 (contract-out [variable-mapping? contract?]
               [string-variable-mapping? contract?]
               [general-pair/c (-> contract? contract? contract?)]))

(module+ test
  (require rackunit))


(define (variable-mapping? dict) (hash/c symbol? any/c))

;;; A string variable mapping is a mapping from variables to strings.
(define (string-variable-mapping? dict) (hash/c symbol? string?))

;;; A contract allowing pairs constructed via cons or via list.
(define (general-pair/c key-contract val-contract)
  (or/c (list/c key-contract val-contract)
        (cons/c key-contract val-contract)))


;;; ===========================
;;; Additional stream utilities
;;; ===========================


;;; Returns the Cartesian product of the given streams.  The result is
;;; a stream whose elements are the elements of the Cartesian product.
;;;
;;; The implementation is inspired from the implementation of
;;; cartesian-product in racket/list.
(define (cartesian-product/stream . ss)
  ;; Cartesian product of two streams, produces an improper pair.
  (define (cp-2 ss1 ss2)
    (for*/stream ([s1 (in-stream ss1)] [s2 (in-stream ss2)]) (cons s1 s2)))
  ;; Fold-right over the list of streams.  The value for the fold is a
  ;; 1-value stream containing the empty list, which makes all the
  ;; lists proper.
  (foldr cp-2 (sequence->stream (in-value (list))) ss))

(module+ test
  (test-case "cartesian-product/stream"
    (check-equal? (stream->list (cartesian-product/stream (in-range 3) (in-range 4 6) '(a b)))
                  '((0 4 a)
                    (0 4 b)
                    (0 5 a)
                    (0 5 b)
                    (1 4 a)
                    (1 4 b)
                    (1 5 a)
                    (1 5 b)
                    (2 4 a)
                    (2 4 b)
                    (2 5 a)
                    (2 5 b)))))


;;; ==================
;;; Boolean operations
;;; ==================

;;; Returns the n-th Cartesian power of the Boolean domain: {0,1}^n.
(define (boolean-power n) (apply cartesian-product (make-list n '(#f #t))))

(module+ test
  (test-case "boolean-power"
    (check-equal? (boolean-power 2) '((#f #f) (#f #t) (#t #f) (#t #t)))))

;;; Like boolean-power, but returns a stream whose elements the
;;; elements of the Cartesian power.
(define (boolean-power/stream n) (apply cartesian-product/stream (make-list n '(#f #t))))

(module+ test
  (test-case "boolean-power/stream"
    (check-equal? (stream->list (boolean-power/stream 2)) '((#f #f) (#f #t) (#t #f) (#t #t)))))

;;; Converts any non-#f value to 1 and #f to 0.
(define (any->01 x) (if x 1 0))

(module+ test
  (test-case "any->01"
    (check-equal? (any->01 #t) 1)
    (check-equal? (any->01 #f) 0)))

;;; Converts 0 to #f and 1 to #t
(define (01->boolean x)
  (case x [(0) #f] [else #t]))

(module+ test
  (test-case "01->boolean"
    (check-equal? (01->boolean 0) #f)
    (check-equal? (01->boolean 1) #t)))