Source file bi_stream.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
open Printf

let error s = failwith ("Bi_stream: " ^ s)

let input_int64 ic =
  match Sys.word_size with
      64 ->
        let n = ref 0 in
        for i = 1 to 8 do
          n := (!n lsl 8) lor (input_byte ic);
        done;
        if !n < 0 then
          error "Corrupted stream: excessive chunk length";
        !n
    | 32 ->
        for i = 1 to 4 do
          if input_byte ic <> 0 then
            error "Chunk length exceeds supported range on this platform"
        done;
        let n = ref 0 in
        for i = 1 to 4 do
          n := (!n lsl 8) lor (input_byte ic);
        done;
        if !n < 0 then
          error "Chunk length exceeds supported range on this platform";
        !n
    | n ->
        error (sprintf "unsupported word size (%i)" n)

let output_int64 oc n =
  match Sys.word_size with
      64 ->
        let n = ref n in
        for i = 1 to 8 do
          output_char oc (char_of_int (!n lsr 56));
          n := !n lsl 8
        done
    | 32 ->
        output_string oc "\000\000\000\000";
        let n = ref n in
        for i = 1 to 4 do
          output_char oc (char_of_int (!n lsr 24));
          n := !n lsl 8
        done
    | n ->
        error (sprintf "unsupported word size (%i)" n)

let rec read_chunk of_string ic =
  match input_char ic with
      '\001' ->
        let len = input_int64 ic in
        if len > Sys.max_string_length then
          error
            (sprintf
               "Corrupted stream: excessive chunk length (%i bytes)" len);
        let s = Bytes.create len in
        really_input ic s 0 len;
        Some (of_string (Bytes.to_string s))

    | '\000' -> None

    | c -> error (sprintf "Corrupted stream: %C" c)


let flatten st =
  let a = ref [| |] in
  let pos = ref 0 in
  let rec next i =
    if !pos >= Array.length !a then (
      match Stream.peek st with
          None -> None
        | Some a' ->
            Stream.junk st;
            a := a';
            pos := 0;
            next i
    )
    else (
      let x = (!a).(!pos) in
      incr pos;
      Some x
    )
  in
  Stream.from next


let read_stream of_string ic =
  flatten (Stream.from (fun i -> read_chunk of_string ic))

let rev_array_of_list l =
  match l with
      [] -> [||]
    | x :: tl ->
        let r = ref tl in
        let len = List.length l in
        let a = Array.make len x in
        for i = len - 2 downto 0 do
          match !r with
              hd :: tl ->
                a.(i) <- hd;
                r := tl;
            | [] -> assert false
        done;
        a

let write_stream ?(chunk_len = 1024) to_string oc st =
  let n = ref 0 in
  let acc = ref [] in
  let flush_chunk () =
    let a = rev_array_of_list !acc in
    acc := [];
    n := 0;
    let s = to_string a in
    output_char oc '\001';
    output_int64 oc (String.length s);
    output_string oc s
  in
  Stream.iter (
    fun x ->
      incr n;
      acc := x :: !acc;
      if !n >= chunk_len then
        flush_chunk ()
  ) st;
  if !n > 0 then
    flush_chunk ();
  output_char oc '\000'


let test l =
  List.iter (fun x -> assert (x >= 0 && x <= 9)) l;
  let to_string a =
    String.concat "" (List.map string_of_int (Array.to_list a))
  in
  let of_string s =
    Array.init (String.length s) (fun i -> int_of_string (String.make 1 s.[i]))
  in
  let st = Stream.of_list l in
  let oc = open_out "test-stream.dat" in
  write_stream ~chunk_len:2 to_string oc st;
  close_out oc;

  let ic = open_in "test-stream.dat" in
  let st' = read_stream of_string ic in
  let l' = ref [] in
  Stream.iter (fun i -> l' := i :: !l') st';
  close_in ic;
  l = List.rev !l'