qcslow2.ml

(*
Install dependencies:

    opam install -y base core ocamlfind

Compile with:

    ocamlfind ocamlopt -g -package base,core,base_quickcheck,base_bigstring -linkpkg qcslow.ml -o qcslow.${OPAMSWITCH}.opt

    ocamlfind ocamlc -package base,core,base_quickcheck,base_bigstring -linkpkg qcslow.ml -o qcslow.byt

Compare:

    OCAMLRUNPARAM=v=0x400 time ./qcslow
    OCAMLRUNPARAM=v=0x400,M=10000 time ./qcslow
*)


open Base
open Core
(*module Bigstring = Base_bigstring*)

module type BS = sig
  type t
  val create : int -> t
  val unsafe_set : t -> int -> char -> unit
  val of_string : string -> t
end

module Main (Bigstring : BS) = struct

  module Column_type = struct
    type 'a t' =
      | Boolean : bool t'
      | Int32 : Int32.t t'
      | Int64 : Int64.t t'
      | Binary : Bigstring.t t'
      | String : Bigstring.t t'

    and 'a t =
      | Required : 'b t' -> 'b t
      | Nullable : 'b t' -> 'b option t
  end

  module Generators = struct
    let gen_column_with_length gen ~length =
      let open Quickcheck.Generator.Let_syntax in
      Let_syntax.bind
        (Int.gen_incl 1 (length / 2 |> Int.max 1))
        ~f:(fun num_unique_values ->
          Let_syntax.bind
            (Quickcheck.Generator.list_with_length num_unique_values gen)
            ~f:(fun unique_values ->
              let num_duplicates_to_include =
                length // num_unique_values |> Float.iround_up_exn
              in
              List.init num_duplicates_to_include ~f:(fun _ -> unique_values)
              |> List.concat
              |> List.sub ~pos:0 ~len:length
              |> List.gen_permutations))
    ;;

    let gen_chunked_column_with_length gen ~length =
      let open Quickcheck.Generator.Let_syntax in
      gen_column_with_length gen ~length
      >>| List.chunks_of ~length:9
      >>| List.map ~f:Array.of_list
    ;;

    module Type = struct
      type 'a t' =
        | Int64 : int64 t'
        | Int32 : int32 t'
        | Utf8 : string t'
        | Byte_array : bytes t'
        | Bool : bool t'

      let element_generator_t' (type a) (t' : a t') : a Quickcheck.Generator.t =
        let open Quickcheck.Generator.Let_syntax in
        let gen_byte_array ~alphabet =
          let alphabet = List.map alphabet ~f:return |> Quickcheck.Generator.union in
          let of_length_between ~min ~max =
            Let_syntax.bind (Int.gen_incl min max) ~f:(fun length ->
              List.gen_with_length length alphabet >>| String.of_char_list)
          in
          Quickcheck.Generator.weighted_union
            [ 1., return ""
            ; 2., of_length_between ~min:1 ~max:8
            ; 2., of_length_between ~min:8 ~max:32
            ; 2., of_length_between ~min:32 ~max:100
            ]
        in
        match t' with
        | Int64 -> Int64.quickcheck_generator
        | Int32 -> Int32.quickcheck_generator
        | Utf8 -> gen_byte_array ~alphabet:[ 'A'; 'B'; 'C' ]
        | Byte_array ->
          gen_byte_array ~alphabet:[ 'A'; 'B'; Char.of_int_exn 0; Char.of_int_exn 0xFF ]
          >>| Bytes.of_string
        | Bool -> Bool.quickcheck_generator
      ;;

      type 'a t =
        | Required : 'a t' -> 'a t
        | Optional : 'a t' -> 'a option t

      type packed = Packed : _ t -> packed

      let packed_constructors =
        [ Packed (Required Int64)
        ; Packed (Optional Int64)
        ; Packed (Required Int32)
        ; Packed (Optional Int32)
        ; Packed (Required Utf8)
        ; Packed (Optional Utf8)
        ; Packed (Required Byte_array)
        ; Packed (Optional Byte_array)
        ; Packed (Required Bool)
        ; Packed (Optional Bool)
        ]
      ;;

      let element_generator (type a) (t : a t) : a Quickcheck.Generator.t =
        let open Quickcheck.Generator.Let_syntax in
        match t with
        | Required ty -> element_generator_t' ty
        | Optional ty ->
          Quickcheck.Generator.weighted_union
            [ 0.9, element_generator_t' ty >>| Option.some; 0.1, return None ]
      ;;

      let quickcheck_generator =
        let no_arg_constructors =
          List.map packed_constructors ~f:Quickcheck.Generator.return
        in
        Quickcheck.Generator.union no_arg_constructors
      ;;
    end

    module Value = struct
      type 'a t = Val : 'a Type.t * 'a -> 'a t
      type packed = Packed : _ t -> packed
    end

    module Column = struct
      type 'a t = Col : 'a Type.t * 'a array list -> 'a t
      type packed = Packed : _ t -> packed

      let to_array (t : packed) : Value.packed array =
        let (Packed (Col (ty, chunks))) = t in
        Array.concat chunks |> Array.map ~f:(fun x -> Value.Packed (Value.Val (ty, x)))
      ;;

      let generate_column_with_type (type a) (ty : a Type.t) ~length
        : a t Quickcheck.Generator.t
        =
        let open Quickcheck.Generator.Let_syntax in
        let gen : a Quickcheck.Generator.t = Type.element_generator ty in
        Let_syntax.map (gen_chunked_column_with_length gen ~length) ~f:(fun col ->
          Col (ty, col))
      ;;

      let quickcheck_generator ~length : packed Quickcheck.Generator.t =
        let open Quickcheck.Generator.Let_syntax in
        Let_syntax.bind Type.quickcheck_generator ~f:(fun (Type.Packed col_type) ->
          Let_syntax.map (generate_column_with_type col_type ~length) ~f:(fun t -> Packed t))
      ;;
    end

    module Table = struct
      type t = { columns : Column.packed array }

      let to_ocaml_rows (t : t) : Value.packed array array =
        let columns = Array.map t.columns ~f:Column.to_array in
        let num_columns = Array.length columns in
        let num_rows = columns.(0) |> Array.length in
        Array.init num_rows ~f:(fun row_idx ->
          Array.init num_columns ~f:(fun col_idx -> columns.(col_idx).(row_idx)))
      ;;

      let quickcheck_generator : t Quickcheck.Generator.t =
        let open Quickcheck.Generator.Let_syntax in
        Let_syntax.bind (Int.gen_incl 1 20) ~f:(fun num_rows ->
          Let_syntax.bind (Column.quickcheck_generator ~length:num_rows) ~f:(fun column ->
            return { columns = [| column |] }))
      ;;
    end
  end

  let stride = 4096
  let createlen =
    let argv = Sys.get_argv () in
    if Array.length argv <= 2 then failwith "usage: qcslow <alloc-size>";
    int_of_string argv.(2)

  module Binary_encoder = struct
    type t = Bigstring.t

    let consume_to_bytes _ = Bytes.create 0
    let create ?(len = createlen) () =
      let result = Bigstring.create len in
      if stride > 0 then begin
        let rec loop i =
          if i < len then begin
            Bigstring.unsafe_set result i ' ';
            loop (i + stride)
          end
        in
        loop 15
      end;
      result
    let[@inline never] write_boolean _ _ = ()
    let[@inline never] write_long _ _ = ()
    let[@inline never] write_int _ _ = ()
    let[@inline never] write_int32 _ _ = ()
    let[@inline never] write_bytes_bigsubstring _ _ = ()
  end

  module Encoder = struct
    type t = Binary_encoder.t

    let create = Binary_encoder.create
    let consume_to_bytes = Binary_encoder.consume_to_bytes

    let write : type a. column_type:a Column_type.t -> t -> a -> unit =
      let write' : type a. column_type:a Column_type.t' -> t -> a -> unit =
        fun ~column_type ->
        match column_type with
        | Boolean -> Binary_encoder.write_boolean
        | Int32 -> Binary_encoder.write_int32
        | Int64 -> Binary_encoder.write_long
        | Binary -> Binary_encoder.write_bytes_bigsubstring
        | String -> Binary_encoder.write_bytes_bigsubstring
      in
      fun ~column_type ->
        match column_type with
        | Required inner_type -> write' ~column_type:inner_type
        | Nullable inner_type ->
          fun t value ->
            (match value with
             | None -> Binary_encoder.write_int t 0
             | Some v ->
               Binary_encoder.write_int t 1;
               write' ~column_type:inner_type t v)
    ;;
  end

  let test_row (row : Generators.Value.packed array) : bytes =
    let encoder = Encoder.create () in
    List.map [ 0 ] ~f:(fun idx -> row.(idx))
    |> Array.of_list
    |> Array.iter ~f:(fun (Generators.Value.Packed value) ->
      match value with
      | Val (Required Generators.Type.Utf8, value) ->
        Encoder.write
          ~column_type:Column_type.(Required String)
          encoder
          (Bigstring.of_string value)
      | Val (Optional Generators.Type.Utf8, value) ->
        Encoder.write
          ~column_type:Column_type.(Nullable String)
          encoder
          (Option.map value ~f:(fun value -> Bigstring.of_string value))
      | Val (Required Generators.Type.Byte_array, value) ->
        Encoder.write
          ~column_type:Column_type.(Required Binary)
          encoder
          (value |> Bytes.to_string |> Bigstring.of_string)
      | Val (Optional Generators.Type.Byte_array, value) ->
        Encoder.write
          ~column_type:Column_type.(Nullable Binary)
          encoder
          (Option.map value ~f:(fun value ->
             value |> Bytes.to_string |> Bigstring.of_string))
      | Val (Required Generators.Type.Int64, value) ->
        Encoder.write ~column_type:Column_type.(Required Int64) encoder value
      | Val (Optional Generators.Type.Int64, value) ->
        Encoder.write ~column_type:Column_type.(Nullable Int64) encoder value
      | Val (Required Generators.Type.Int32, value) ->
        Encoder.write ~column_type:Column_type.(Required Int32) encoder value
      | Val (Optional Generators.Type.Int32, value) ->
        Encoder.write ~column_type:Column_type.(Nullable Int32) encoder value
      | Val (Required Generators.Type.Bool, value) ->
        Encoder.write ~column_type:Column_type.(Required Boolean) encoder value
      | Val (Optional Generators.Type.Bool, value) ->
        Encoder.write ~column_type:Column_type.(Nullable Boolean) encoder value);
    Encoder.consume_to_bytes encoder
  ;;

  let test_encoder table =
    let rows = Generators.Table.to_ocaml_rows table in
    let output = Array.map rows ~f:test_row in
    ignore (Sys.opaque_identity output : _);
    ()
  ;;

  let main () =
    let ballast =
      let rec loop i acc = if i <= 0 then acc else loop (i-1) (i::acc) in
      loop 1_000_000 []
    in
    Base_quickcheck.Test.run_exn
      ~config:{ Base_quickcheck.Test.default_config with test_count = 1000 }
      ~examples:[]
      (module struct
        type t = Generators.Table.t

        let sexp_of_t _t = Sexp.List []
        let quickcheck_generator = Generators.Table.quickcheck_generator
        let quickcheck_shrinker = Base_quickcheck.Shrinker.atomic
      end)
      ~f:test_encoder;
    Printf.printf "%d\n" (List.hd_exn ballast)
  ;;
end

module BS = struct
  include Base_bigstring
  let of_string x = Base_bigstring.of_string x
  let create n = Base_bigstring.create n
end

module WithBytes = Main (Bytes)
module WithBigstring = Main (BS)

;;
let argv = Sys.get_argv () in
if Array.length argv > 1 && Stdlib.(=) argv.(1) "by" then
  WithBytes.main ()
else
  WithBigstring.main ()