#lang racket/base
;; "varints" from Google Protocol Buffers,
;; https://developers.google.com/protocol-buffers/docs/encoding#varints
;;
;; "Each byte in a varint, except the last byte, has the most
;;  significant bit (msb) set – this indicates that there are further
;;  bytes to come. The lower 7 bits of each byte are used to store the
;;  two's complement representation of the number in groups of 7 bits,
;;  least significant group first."

(provide write-varint
         read-varint
         encode-varint
         decode-varint)

(require racket/port)

(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))))

(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))]))))

(define (encode-varint v)
  (call-with-output-bytes (lambda (p) (write-varint v p))))

(define (decode-varint bs ks kf)
  ((call-with-input-bytes bs (lambda (p)
                               (define v (read-varint p))
                               (cond [(eof-object? v) (lambda () (kf #t))]
                                     [else (define rest (port->bytes p))
                                           (lambda () (ks v rest))])))))

(module+ test
  (require rackunit)

  (check-equal? (encode-varint 0) (bytes 0))
  (check-equal? (encode-varint 1) (bytes 1))
  (check-equal? (encode-varint 127) (bytes 127))
  (check-equal? (encode-varint 128) (bytes 128 1))
  (check-equal? (encode-varint 255) (bytes 255 1))
  (check-equal? (encode-varint 256) (bytes 128 2))
  (check-equal? (encode-varint 300) (bytes #b10101100 #b00000010))
  (check-equal? (encode-varint 1000000000) (bytes 128 148 235 220 3))

  (define (ks* v rest) (list v rest))
  (define (kf* [short? #f]) (if short? 'short (void)))

  (check-equal? (decode-varint (bytes) ks* kf*) 'short)
  (check-equal? (decode-varint (bytes 0) ks* kf*) (list 0 (bytes)))
  (check-equal? (decode-varint (bytes 0 99) ks* kf*) (list 0 (bytes 99)))
  (check-equal? (decode-varint (bytes 1) ks* kf*) (list 1 (bytes)))
  (check-equal? (decode-varint (bytes 127) ks* kf*) (list 127 (bytes)))
  (check-equal? (decode-varint (bytes 128) ks* kf*) 'short)
  (check-equal? (decode-varint (bytes 128 1) ks* kf*) (list 128 (bytes)))
  (check-equal? (decode-varint (bytes 128 1 99) ks* kf*) (list 128 (bytes 99)))
  (check-equal? (decode-varint (bytes 255 1) ks* kf*) (list 255 (bytes)))
  (check-equal? (decode-varint (bytes 128 2) ks* kf*) (list 256 (bytes)))
  (check-equal? (decode-varint (bytes #b10101100 #b00000010) ks* kf*) (list 300 (bytes)))
  (check-equal? (decode-varint (bytes 128 148 235 220 3) ks* kf*) (list 1000000000 (bytes)))
  (check-equal? (decode-varint (bytes 128 148 235 220 3 99) ks* kf*) (list 1000000000 (bytes 99))))