#lang racket
;; Jelly, a very shaky implementation of Preserves - intended to
;; demonstrate a minimal implementation of Preserves binary I/O,
;; without error-checking or configurability etc.

(provide (all-defined-out))

;;---------------------------------------------------------------------------
;; Representing values

(struct record (label fields) #:transparent)
(struct float (value) #:transparent) ;; a marker for single-precision I/O
(struct annotated (annotations item) #:transparent)
(struct embedded (value) #:transparent)

;;---------------------------------------------------------------------------
;; Reader

(define (read-preserve/binary [in-port (current-input-port)])
  (let/ec return

    (define (next)
      (match (next-byte)
        [#x80 #f]
        [#x81 #t]
        [#x82 (float (floating-point-bytes->real (next-bytes 4) #t 0 4))]
        [#x83 (floating-point-bytes->real (next-bytes 8) #t 0 8)]
        [#x84 '#:end]
        [#x85 (let ((a (next)))
                (match (next)
                  [(annotated as i) (annotated (cons a as) i)]
                  [i (annotated (list a) i)]))]
        [#x86 (embedded (next))]
        [(? (between #x90 #x9C) v) (- v #x90)]
        [(? (between #x9D #x9F) v) (- v #xA0)]
        [(? (between #xA0 #xAF) v) (next-integer (- v #xA0 -1))]
        [#xB0 (next-integer (next-varint))]
        [#xB1 (bytes->string/utf-8 (next-bytes (next-varint)))]
        [#xB2 (next-bytes (next-varint))]
        [#xB3 (string->symbol (bytes->string/utf-8 (next-bytes (next-varint))))]
        [#xB4 (apply (lambda (label . fields) (record label fields)) (next-items))]
        [#xB5 (next-items)]
        [#xB6 (list->set (next-items))]
        [#xB7 (apply hash (next-items))]))

    (define (next-items) (match (next) ['#:end '()] [v (cons v (next-items))]))

    (define (eof-guard v) (if (eof-object? v) (return eof) v))

    (define (next-byte) (eof-guard (read-byte in-port)))

    (define (next-bytes n)
      (define bs (eof-guard (read-bytes n in-port)))
      (if (< (bytes-length bs) n) (return eof) bs))

    (define (next-varint) (eof-guard (read-varint in-port)))

    (define (next-integer n)
      (define acc0 (next-byte))
      (define acc (if (< acc0 128) acc0 (- acc0 256)))
      (for/fold [(acc acc)] [(n (in-range (- n 1)))] (+ (* acc 256) (next-byte))))

    (next)))

(define ((between lo hi) v) (<= lo v hi))

(define (read-varint in-port)
  (let/ec return
    (let loop ()
      (define b (read-byte in-port))
      (cond [(eof-object? b) (return b)]
            [(< b 128) b]
            [else (+ (* (loop) 128) (- b 128))]))))

;;---------------------------------------------------------------------------
;; Writer

(define (write-preserve/binary v [out-port (current-output-port)])
  (define (output v)
    (match v
      [#f (write-byte #x80 out-port)]
      [#t (write-byte #x81 out-port)]
      [(float v) (write-byte #x82 out-port) (output-bytes (real->floating-point-bytes v 4 #t))]
      [(? flonum?) (write-byte #x83 out-port) (output-bytes (real->floating-point-bytes v 8 #t))]

      [(annotated as v)
       (for [(a (in-list as))] (write-byte #x85 out-port) (output a))
       (output v)]

      [(embedded v) (write-byte #x86 out-port) (output v)]

      [(? integer?)
       (cond [(<= -3 v -1) (write-byte (+ v #xA0) out-port)]
             [(<= 0 v 12)  (write-byte (+ v #x90) out-port)]
             [else (define raw-bit-count (+ (integer-length v) 1)) ;; at least one sign bit
                   (define byte-count (quotient (+ raw-bit-count 7) 8))
                   (if (<= byte-count 16)
                       (write-byte (+ byte-count #xA0 -1) out-port)
                       (begin (write-byte #xB0 out-port)
                              (write-varint byte-count out-port)))
                   (for [(shift (in-range (* byte-count 8) 0 -8))]
                     (write-byte (bitwise-bit-field v (- shift 8) shift) out-port))])]

      [(? string?) (count-bytes 1 (string->bytes/utf-8 v))]
      [(? bytes?)  (count-bytes 2 v)]
      [(? symbol?) (count-bytes 3 (string->bytes/utf-8 (symbol->string v)))]

      [(record label fields) (with-seq 4 (output label) (for-each output fields))]
      [(? list?)             (with-seq 5 (for-each output v))]
      [(? set?)              (with-seq 6 (output-set v))]
      [(? hash?)             (with-seq 7 (output-hash v))]

      [_ (error 'write-preserve/binary "Invalid value: ~v" v)]))

  (define (output-bytes bs) (write-bytes bs out-port))

  (define-syntax-rule (with-seq tag body ...)
    (begin (write-byte (+ tag #xB0) out-port)
           body ...
           (write-byte #x84 out-port)))

  (define (count-bytes tag bs)
    (write-byte (+ tag #xB0) out-port)
    (write-varint (bytes-length bs) out-port)
    (output-bytes bs))

  (define (encode v) (call-with-output-bytes (lambda (p) (write-preserve/binary v p))))

  (define (output-set v)
    (for-each output-bytes (sort (for/list [(e (in-set v))] (encode e)) bytes<?)))

  (define (output-hash d)
    (define encoded-entries (for/list [((k v) (in-hash d))] (cons (encode k) (encode v))))
    (for-each output-bytes (flatten (sort encoded-entries bytes<? #:key car))))

  (output v))

(define (write-varint v out-port)
  (if (< v 128)
      (write-byte v out-port)
      (begin (write-byte (+ 128 (modulo v 128)) out-port)
             (write-varint (quotient v 128) out-port))))