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
149
150
151
152
153
154
155
156
|
open! Import
open! Map
open! Int
module Lazy_apply = struct
module T = struct
type 'a t = { compute : unit -> 'a } [@@unboxed]
let run t = t.compute ()
let return x = { compute = (fun () -> x) }
let map x ~f = { compute = (fun () -> f (run x)) }
let both x y = { compute = (fun () -> run x, run y) }
let map2 a b ~f = map (both a b) ~f:(fun (x, y) -> f x y)
let map = `Custom map
end
include T
include Applicative.Make_using_map2 (T)
let of_thunk f = { compute = (fun () -> run (f ())) }
end
module Lazy_map = Map.Make_applicative_traversals (Lazy_apply)
let%expect_test "mapi correctness check" =
let map = Map.of_alist_exn (module Int) (List.init 100 ~f:(fun x -> x, x)) in
let f ~key:_ ~data = data + 1 in
let test_output =
Lazy_map.mapi map ~f:(fun ~key ~data -> Lazy_apply.return (f ~key ~data))
|> Lazy_apply.run
in
let reference_output = Map.mapi map ~f in
require [%here] (Map.equal Int.equal test_output reference_output);
require [%here] (Map.invariants test_output);
[%expect {| |}]
;;
let%expect_test "filter_mapi correctness check" =
let map = Map.of_alist_exn (module Int) (List.init 1000 ~f:(fun x -> x, x)) in
let f ~key:_ ~data = if data % 50 > 10 then None else Some data in
let test_output =
Lazy_map.filter_mapi map ~f:(fun ~key ~data -> Lazy_apply.return (f ~key ~data))
|> Lazy_apply.run
in
let reference_output = Map.filter_mapi map ~f in
require [%here] (Map.equal Int.equal test_output reference_output);
require [%here] (Map.invariants test_output);
[%expect {| |}]
;;
module Step_applicative = struct
module M = struct
type 'a t = { compute : steps:int -> ('a * int, 'a t) Either.t }
let return x = { compute = (fun ~steps -> First (x, steps)) }
let step x =
let rec t =
{ compute = (fun ~steps -> if steps > 0 then First (x, steps - 1) else Second t) }
in
t
;;
let internal_map x ~f =
let rec fn t =
{ compute =
(fun ~steps ->
match t.compute ~steps with
| First (x, steps) -> First (f x, steps)
| Second t -> Second (fn t))
}
in
fn x
;;
let map2 a b ~f =
let rec fn a =
{ compute =
(fun ~steps ->
match a.compute ~steps with
| First (x, steps) -> (internal_map b ~f:(fun y -> f x y)).compute ~steps
| Second t -> Second (fn t))
}
in
fn a
;;
let map = `Custom internal_map
let of_thunk f =
{ compute =
(fun ~steps ->
let t = f () in
t.compute ~steps)
}
;;
end
include M
include Applicative.Make_using_map2 (M)
end
module Step_map = Map.Make_applicative_traversals (Step_applicative)
let%expect_test "mapi lazy check" =
let map = Map.of_alist_exn (module Int) (List.init 10 ~f:(fun x -> x, x)) in
let f ~key:_ ~data = data * 2 in
(* transform the map, expect no output yet *)
let step_computation =
Step_map.mapi map ~f:(fun ~key ~data ->
Step_applicative.of_thunk (fun () ->
print_s [%message (key : int) (data : int)];
Step_applicative.step (f ~key ~data)))
in
[%expect {| |}];
(* take a few steps, expect some but not all output *)
let more_computation =
match step_computation.compute ~steps:3 with
| First _ -> assert false
| Second c -> c
in
[%expect
{|
((key 0)
(data 0))
((key 1)
(data 1))
((key 2)
(data 2))
((key 3)
(data 3))
|}];
(* take more than enough steps to finish, expect the rest of the output *)
let test_output =
match more_computation.compute ~steps:100 with
| First (r, _) -> r
| Second _ -> assert false
in
let reference_output = Map.mapi map ~f in
require [%here] (Map.equal Int.equal test_output reference_output);
[%expect
{|
((key 4)
(data 4))
((key 5)
(data 5))
((key 6)
(data 6))
((key 7)
(data 7))
((key 8)
(data 8))
((key 9)
(data 9))
|}]
;;
|
