Source file path_trie.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
type 'a t = {
  value : 'a option;
  children : 'a t String_map.t;
}

let empty : 'a t =
  {
    value = None;
    children = String_map.empty;
  }

let add (path : Abs_path.t) (v : 'a) (t : 'a t) : 'a t =
  let rec aux t parts =
    match parts with
    | [] -> { t with value = Some v }
    | x :: xs ->
      let children =
        String_map.find_opt x t.children
        |> Option.value ~default:empty
        |> (fun sub_trie ->
            String_map.add x (aux sub_trie xs) t.children
          )
      in
      { t with children }
  in
  aux t (Abs_path.to_parts path)

let remove (path : Abs_path.t) (t : 'a t) : 'a t =
  let rec aux t parts =
    match parts with
    | [] -> { t with value = None }
    | x :: xs -> (
        match String_map.find_opt x t.children with
        | None -> t
        | Some sub_trie ->
          let children =
            String_map.add x (aux sub_trie xs) t.children
          in
          { t with children }
      )
  in
  aux t (Abs_path.to_parts path)

let find (path : Abs_path.t) (t : 'a t) : 'a option =
  let rec aux t parts =
    match parts with
    | [] -> t.value
    | x :: xs ->
      match String_map.find_opt x t.children with
      | None -> None
      | Some t -> aux t xs
  in
  aux t (Abs_path.to_parts path)

let find_exn path t =
  match find path t with
  | None -> invalid_arg "find_exn: Path does not exist"
  | Some x -> x

let to_seq (t : 'a t) : (Abs_path.t * 'a) Seq.t =
  let rec aux (t : 'a t) : (string list * 'a) Seq.t =
    let sub_tries_seq =
      String_map.to_seq t.children
      |> Seq.flat_map (fun (k, sub_trie) ->
          Seq.map (fun (l, v) ->
              (k :: l, v)
            )
            (aux sub_trie)
        )
    in
    match t.value with
    | None -> sub_tries_seq
    | Some v ->
      CCSeq.cons ([], v) sub_tries_seq
  in
  aux t
  |> Seq.map (fun (l, v) ->
      (Abs_path.of_parts_exn l, v)
    )

let of_seq (s : (Abs_path.t * 'a) Seq.t) : 'a t =
  Seq.fold_left (fun acc (p, x) ->
      add p x acc
    )
    empty
    s

let is_empty t =
  match to_seq t () with
  | Seq.Nil -> true
  | _ -> false

let merge
    (type a b c)
    (f : Abs_path.t -> a option -> b option -> c option)
    (t1 : a t)
    (t2 : b t)
  : c t =
  let s1 = to_seq t1 |> Seq.map fst in
  let s2 = to_seq t2 |> Seq.map fst in
  let keys = CCSeq.append s1 s2
             |> Seq.fold_left (fun acc p ->
                 Abs_path_set.add p acc
               )
               Abs_path_set.empty
  in
  Abs_path_set.to_seq keys
  |> Seq.fold_left (fun t p ->
      match f p (find p t1) (find p t2) with
      | None -> t
      | Some x -> add p x t
    )
    empty

let union
    (type a)
    (f : Abs_path.t -> a -> a -> a option)
    (t1 : a t)
    (t2 : a t)
  : a t =
  merge (fun p x y ->
      match x, y with
      | None, None -> None
      | Some v, None -> Some v
      | None, Some v -> Some v
      | Some x, Some y -> f p x y
    ) t1 t2

let equal (f : 'a -> 'a -> bool) t1 t2 =
  CCSeq.equal (fun (p1, v1) (p2, v2) ->
      Abs_path.equal p1 p2 && f v1 v2
    )
    (to_seq t1) (to_seq t2)