Source file Belt_MutableMapInt.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
149
150
151
152
153
154
155
156
157
158
159
module I = Belt_internalMapInt

type key = int

module N = Belt_internalAVLtree
module A = Belt_Array

include (
  struct
    type 'a t = { mutable data : 'a I.t }

    let t : data:'a I.t -> 'a t = fun ~data -> { data }
    let dataSet : 'a t -> 'a I.t -> unit = fun o v -> o.data <- v
    let data : 'a t -> 'a I.t = fun o -> o.data
  end :
    sig
      type 'a t

      val t : data:'a I.t -> 'a t
      val dataSet : 'a t -> 'a I.t -> unit
      val data : 'a t -> 'a I.t
    end)

let make () = t ~data:N.empty
let isEmpty m = N.isEmpty (data m)
let clear m = dataSet m N.empty
let singleton k v = t ~data:(N.singleton k v)
let minKeyUndefined m = N.minKeyUndefined (data m)
let minKey m = N.minKey (data m)
let maxKeyUndefined m = N.maxKeyUndefined (data m)
let maxKey m = N.maxKey (data m)
let minimum m = N.minimum (data m)
let minUndefined m = N.minUndefined (data m)
let maximum m = N.maximum (data m)
let maxUndefined m = N.maxUndefined (data m)

let set (m : _ t) k v =
  let old_data = data m in
  let v = I.addMutate old_data k v in
  if v != old_data then dataSet m v

let forEachU d f = N.forEachU (data d) f
let forEach d f = forEachU d (fun a b -> f a b)
let mapU d f = t ~data:(N.mapU (data d) f)
let map d f = mapU d (fun a -> f a)
let mapWithKeyU d f = t ~data:(N.mapWithKeyU (data d) f)
let mapWithKey d f = mapWithKeyU d (fun a b -> f a b)
let reduceU d acc f = N.reduceU (data d) acc f
let reduce d acc f = reduceU d acc (fun a b c -> f a b c)
let everyU d f = N.everyU (data d) f
let every d f = everyU d (fun a b -> f a b)
let someU d f = N.someU (data d) f
let some d f = someU d (fun a b -> f a b)
let size d = N.size (data d)
let toList d = N.toList (data d)
let toArray d = N.toArray (data d)
let keysToArray d = N.keysToArray (data d)
let valuesToArray d = N.valuesToArray (data d)
let checkInvariantInternal d = N.checkInvariantInternal (data d)
let has d v = I.has (data d) v

let rec removeMutateAux nt (x : key) =
  let k = N.key nt in
  if x = k then (
    let l, r =
      let open N in
      (left nt, right nt)
    in
    match
      let open N in
      (toOpt l, toOpt r)
    with
    | None, _ -> r
    | _, None -> l
    | _, Some nr ->
        N.rightSet nt (N.removeMinAuxWithRootMutate nt nr);
        N.return (N.balMutate nt))
  else if x < k then (
    match N.toOpt (N.left nt) with
    | None -> N.return nt
    | Some l ->
        N.leftSet nt (removeMutateAux l x);
        N.return (N.balMutate nt))
  else
    match N.toOpt (N.right nt) with
    | None -> N.return nt
    | Some r ->
        N.rightSet nt (removeMutateAux r x);
        N.return (N.balMutate nt)

let remove d v =
  let oldRoot = data d in
  match N.toOpt oldRoot with
  | None -> ()
  | Some root ->
      let newRoot = removeMutateAux root v in
      if newRoot != oldRoot then dataSet d newRoot

let rec updateDone t (x : key) f =
  match N.toOpt t with
  | None -> ( match f None with Some data -> N.singleton x data | None -> t)
  | Some nt ->
      let k = N.key nt in
      if k = x then (
        match f (Some (N.value nt)) with
        | None -> (
            let l, r = (N.left nt, N.right nt) in
            match (N.toOpt l, N.toOpt r) with
            | None, _ -> r
            | _, None -> l
            | _, Some nr ->
                N.rightSet nt (N.removeMinAuxWithRootMutate nt nr);
                N.return (N.balMutate nt))
        | Some data ->
            N.valueSet nt data;
            N.return nt)
      else
        let l, r =
          let open N in
          (left nt, right nt)
        in
        if x < k then
          let ll = updateDone l x f in
          N.leftSet nt ll
        else N.rightSet nt (updateDone r x f);
        N.return (N.balMutate nt)

let updateU t x f =
  let oldRoot = data t in
  let newRoot = updateDone oldRoot x f in
  if newRoot != oldRoot then dataSet t newRoot

let update t x f = updateU t x (fun a -> f a)

let rec removeArrayMutateAux t xs i len =
  if i < len then
    let ele = A.getUnsafe xs i in
    let u = removeMutateAux t ele in
    match N.toOpt u with None -> N.empty | Some t -> removeArrayMutateAux t xs (i + 1) len
  else N.return t

let removeMany (type key id) (d : _ t) xs =
  let oldRoot = data d in
  match N.toOpt oldRoot with
  | None -> ()
  | Some nt ->
      let len = A.length xs in
      let newRoot = removeArrayMutateAux nt xs 0 len in
      if newRoot != oldRoot then dataSet d newRoot

let fromArray xs = t ~data:(I.fromArray xs)
let cmpU d0 d1 f = I.cmpU (data d0) (data d1) f
let cmp d0 d1 f = cmpU d0 d1 (fun a b -> f a b)
let eqU d0 d1 f = I.eqU (data d0) (data d1) f
let eq d0 d1 f = eqU d0 d1 (fun a b -> f a b)
let get d x = I.get (data d) x
let getUndefined d x = I.getUndefined (data d) x
let getWithDefault d x def = I.getWithDefault (data d) x def
let getExn d x = I.getExn (data d) x