path: root/stdlib/source/lux/data/binary.lux

(.module:
  [lux (#- i64)
   ["@" target]
   ["." host]
   [abstract
    [monad (#+ do)]
    [equivalence (#+ Equivalence)]
    [monoid (#+ Monoid)]]
   [control
    ["." try (#+ Try)]
    ["." exception (#+ exception:)]]
   [data
    ["." maybe]
    [text
     ["%" format (#+ format)]]
    [collection
     ["." array]]]
   [math
    [number (#+ hex)
     ["n" nat]
     ["f" frac]
     ["." i64]]]])

(exception: #export (index_out_of_bounds {size Nat} {index Nat})
  (exception.report
   ["Size" (%.nat size)]
   ["Index" (%.nat index)]))

(template [<name>]
  [(exception: #export (<name> {size Nat} {from Nat} {to Nat})
     (exception.report
      ["Size" (%.nat size)]
      ["From" (%.nat from)]
      ["To" (%.nat to)]))]

  [slice_out_of_bounds]
  [inverted_slice]
  )

(with_expansions [<jvm> (as_is (type: #export Binary (host.type [byte]))

                               (host.import: java/lang/Object)
                               
                               (host.import: java/lang/System
                                 ["#::."
                                  (#static arraycopy [java/lang/Object int java/lang/Object int int] #try void)])

                               (host.import: java/util/Arrays
                                 ["#::."
                                  (#static copyOfRange [[byte] int int] [byte])
                                  (#static equals [[byte] [byte]] boolean)])

                               (def: byte_mask
                                 I64
                                 (|> i64.bits_per_byte i64.mask .i64))

                               (def: i64
                                 (-> (primitive "java.lang.Byte") I64)
                                 (|>> host.byte_to_long (:coerce I64) (i64.and ..byte_mask)))

                               (def: byte
                                 (-> (I64 Any) (primitive "java.lang.Byte"))
                                 (for {@.old
                                       (|>> .int host.long_to_byte)

                                       @.jvm
                                       (|>> .int (:coerce (primitive "java.lang.Long")) host.long_to_byte)})))]
  (for {@.old (as_is <jvm>)
        @.jvm (as_is <jvm>)

        @.js
        (as_is (host.import: ArrayBuffer
                 (new [host.Number]))
               
               (host.import: Uint8Array
                 (new [ArrayBuffer])
                 (length host.Number))
               
               (type: #export Binary
                 Uint8Array))

        @.python
        (type: #export Binary
          (primitive "bytearray"))

        @.scheme
        (as_is (type: #export Binary
                 (primitive "bytevector"))
               
               (host.import: (make-bytevector [Nat] Binary))
               (host.import: (bytevector-u8-ref [Binary Nat] I64))
               (host.import: (bytevector-u8-set! [Binary Nat (I64 Any)] Any))
               (host.import: (bytevector-length [Binary] Nat)))}

       ## Default
       (type: #export Binary
         (array.Array (I64 Any)))))

(template: (!size binary)
  (for {@.old (host.array_length binary)
        @.jvm (host.array_length binary)

        @.js
        (|> binary
            Uint8Array::length
            f.nat)

        @.python
        (|> binary
            (:coerce (array.Array (I64 Any)))
            "python array length")

        @.scheme
        (..bytevector-length [binary])}

       ## Default
       (array.size binary)))

(template: (!read idx binary)
  (for {@.old (..i64 (host.array_read idx binary))
        @.jvm (..i64 (host.array_read idx binary))

        @.js
        (|> binary
            (: ..Binary)
            (:coerce (array.Array .Frac))
            ("js array read" idx)
            f.nat
            .i64)

        @.python
        (|> binary
            (:coerce (array.Array .I64))
            ("python array read" idx))

        @.scheme
        (..bytevector-u8-ref [binary idx])}

       ## Default
       (|> binary
           (array.read idx)
           (maybe.default (: (I64 Any) 0))
           (:coerce I64))))

(template: (!!write <byte_type> <post> <write> idx value binary)
  (|> binary
      (: ..Binary)
      (:coerce (array.Array <byte_type>))
      (<write> idx (|> value .nat (n.% (hex "100")) <post>))
      (:coerce ..Binary)))

(template: (!write idx value binary)
  (for {@.old (host.array_write idx (..byte value) binary)
        @.jvm (host.array_write idx (..byte value) binary)

        @.js (!!write .Frac n.frac "js array write" idx value binary)
        @.python (!!write (I64 Any) (:coerce (I64 Any)) "python array write" idx value binary)
        @.scheme (exec (..bytevector-u8-set! [binary idx value])
                   binary)}

       ## Default
       (array.write! idx (|> value .nat (n.% (hex "100"))) binary)))

(def: #export size
  (-> Binary Nat)
  (|>> !size))

(def: #export create
  (-> Nat Binary)
  (for {@.old (|>> (host.array byte))
        @.jvm (|>> (host.array byte))

        @.js
        (|>> n.frac ArrayBuffer::new Uint8Array::new)

        @.python
        (|>> ("python apply" (:coerce host.Function ("python constant" "bytearray")))
             (:coerce Binary))

        @.scheme
        (|>> ..make-bytevector)}

       ## Default
       array.new))

(def: #export (fold f init binary)
  (All [a] (-> (-> I64 a a) a Binary a))
  (let [size (..!size binary)]
    (loop [idx 0
           output init]
      (if (n.< size idx)
        (recur (inc idx) (f (!read idx binary) output))
        output))))

(def: #export (read/8 idx binary)
  (-> Nat Binary (Try I64))
  (if (n.< (..!size binary) idx)
    (#try.Success (!read idx binary))
    (exception.throw ..index_out_of_bounds [(..!size binary) idx])))

(def: #export (read/16 idx binary)
  (-> Nat Binary (Try I64))
  (if (n.< (..!size binary) (n.+ 1 idx))
    (#try.Success ($_ i64.or
                      (i64.left_shift 8 (!read idx binary))
                      (!read (n.+ 1 idx) binary)))
    (exception.throw ..index_out_of_bounds [(..!size binary) idx])))

(def: #export (read/32 idx binary)
  (-> Nat Binary (Try I64))
  (if (n.< (..!size binary) (n.+ 3 idx))
    (#try.Success ($_ i64.or
                      (i64.left_shift 24 (!read idx binary))
                      (i64.left_shift 16 (!read (n.+ 1 idx) binary))
                      (i64.left_shift 8 (!read (n.+ 2 idx) binary))
                      (!read (n.+ 3 idx) binary)))
    (exception.throw ..index_out_of_bounds [(..!size binary) idx])))

(def: #export (read/64 idx binary)
  (-> Nat Binary (Try I64))
  (if (n.< (..!size binary) (n.+ 7 idx))
    (#try.Success ($_ i64.or
                      (i64.left_shift 56 (!read idx binary))
                      (i64.left_shift 48 (!read (n.+ 1 idx) binary))
                      (i64.left_shift 40 (!read (n.+ 2 idx) binary))
                      (i64.left_shift 32 (!read (n.+ 3 idx) binary))
                      (i64.left_shift 24 (!read (n.+ 4 idx) binary))
                      (i64.left_shift 16 (!read (n.+ 5 idx) binary))
                      (i64.left_shift 8 (!read (n.+ 6 idx) binary))
                      (!read (n.+ 7 idx) binary)))
    (exception.throw ..index_out_of_bounds [(..!size binary) idx])))

(def: #export (write/8 idx value binary)
  (-> Nat (I64 Any) Binary (Try Binary))
  (if (n.< (..!size binary) idx)
    (#try.Success (|> binary
                      (!write idx value)))
    (exception.throw ..index_out_of_bounds [(..!size binary) idx])))

(def: #export (write/16 idx value binary)
  (-> Nat (I64 Any) Binary (Try Binary))
  (if (n.< (..!size binary) (n.+ 1 idx))
    (#try.Success (|> binary
                      (!write idx (i64.right_shift 8 value))
                      (!write (n.+ 1 idx) value)))
    (exception.throw ..index_out_of_bounds [(..!size binary) idx])))

(def: #export (write/32 idx value binary)
  (-> Nat (I64 Any) Binary (Try Binary))
  (if (n.< (..!size binary) (n.+ 3 idx))
    (#try.Success (|> binary
                      (!write idx (i64.right_shift 24 value))
                      (!write (n.+ 1 idx) (i64.right_shift 16 value))
                      (!write (n.+ 2 idx) (i64.right_shift 8 value))
                      (!write (n.+ 3 idx) value)))
    (exception.throw ..index_out_of_bounds [(..!size binary) idx])))

(def: #export (write/64 idx value binary)
  (-> Nat (I64 Any) Binary (Try Binary))
  (if (n.< (..!size binary) (n.+ 7 idx))
    (for {@.scheme (let [write_high (|>> (!write idx (i64.right_shift 56 value))
                                         (!write (n.+ 1 idx) (i64.right_shift 48 value))
                                         (!write (n.+ 2 idx) (i64.right_shift 40 value))
                                         (!write (n.+ 3 idx) (i64.right_shift 32 value)))
                         write_low (|>> (!write (n.+ 4 idx) (i64.right_shift 24 value))
                                        (!write (n.+ 5 idx) (i64.right_shift 16 value))
                                        (!write (n.+ 6 idx) (i64.right_shift 8 value))
                                        (!write (n.+ 7 idx) value))]
                     (|> binary write_high write_low #try.Success))}
         (#try.Success (|> binary
                           (!write idx (i64.right_shift 56 value))
                           (!write (n.+ 1 idx) (i64.right_shift 48 value))
                           (!write (n.+ 2 idx) (i64.right_shift 40 value))
                           (!write (n.+ 3 idx) (i64.right_shift 32 value))
                           (!write (n.+ 4 idx) (i64.right_shift 24 value))
                           (!write (n.+ 5 idx) (i64.right_shift 16 value))
                           (!write (n.+ 6 idx) (i64.right_shift 8 value))
                           (!write (n.+ 7 idx) value))))
    (exception.throw ..index_out_of_bounds [(..!size binary) idx])))

(structure: #export equivalence
  (Equivalence Binary)
  
  (def: (= reference sample)
    (with_expansions [<jvm> (java/util/Arrays::equals reference sample)]
     (for {@.old <jvm>
           @.jvm <jvm>}
          (let [limit (!size reference)]
            (and (n.= limit
                      (!size sample))
                 (loop [idx 0]
                   (if (n.< limit idx)
                     (and (n.= (!read idx reference)
                               (!read idx sample))
                          (recur (inc idx)))
                     true))))))))

(for {@.old (as_is)
      @.jvm (as_is)}

     ## Default
     (exception: #export (cannot_copy_bytes {bytes Nat}
                                            {source_input Nat}
                                            {target_output Nat})
       (exception.report
        ["Bytes" (%.nat bytes)]
        ["Source input space" (%.nat source_input)]
        ["Target output space" (%.nat target_output)])))

(def: #export (copy bytes source_offset source target_offset target)
  (-> Nat Nat Binary Nat Binary (Try Binary))
  (with_expansions [<jvm> (as_is (do try.monad
                                   [_ (java/lang/System::arraycopy source (.int source_offset) target (.int target_offset) (.int bytes))]
                                   (wrap target)))]
    (for {@.old <jvm>
          @.jvm <jvm>}
         
         ## Default
         (let [source_input (n.- source_offset (!size source))
               target_output (n.- target_offset (!size target))]
           (if (n.<= source_input bytes)
             (loop [idx 0]
               (if (n.< bytes idx)
                 (exec (!write (n.+ target_offset idx)
                               (!read (n.+ source_offset idx) source)
                               target)
                   (recur (inc idx)))
                 (#try.Success target)))
             (exception.throw ..cannot_copy_bytes [bytes source_input target_output]))))))

(def: #export (slice from to binary)
  (-> Nat Nat Binary (Try Binary))
  (let [size (..!size binary)]
    (if (n.<= to from)
      (if (and (n.< size from)
               (n.< size to))
        (with_expansions [<jvm> (as_is (#try.Success (java/util/Arrays::copyOfRange binary (.int from) (.int (inc to)))))]
          (for {@.old <jvm>
                @.jvm <jvm>}
               
               ## Default
               (let [how_many (inc (n.- from to))]
                 (..copy how_many from binary 0 (..create how_many)))))
        (exception.throw ..slice_out_of_bounds [size from to]))
      (exception.throw ..inverted_slice [size from to]))))

(def: #export (drop from binary)
  (-> Nat Binary Binary)
  (case from
    0 binary
    _ (case (..slice from (dec (..!size binary)) binary)
        (#try.Success slice)
        slice
        
        (#try.Failure _)
        (..create 0))))

(structure: #export monoid
  (Monoid Binary)

  (def: identity
    (..create 0))

  (def: (compose left right)
    (let [sizeL (!size left)
          sizeR (!size right)
          output (..create (n.+ sizeL sizeR))]
      (exec
        (..copy sizeL 0 left 0 output)
        (..copy sizeR 0 right sizeL output)
        output))))