Remove fns

src/rtree.ml

@@ -129,6 +129,61 @@ module Make (E : Envelope) (V : Value with type envelope = E.t) = struct
     | a, b -> { max_node_load = t.max_node_load; tree = Node [ a; b ] }
   (* root split *)
 
+  let rec remove_eq' eq = function
+    | Node ns ->
+        let opts, ns' =
+          List.map
+            (fun (e, t) ->
+              let opt, t' = remove_eq' eq t in
+              (opt, (e, t')))
+            ns
+          |> List.split
+        in
+        let opts = List.concat opts in
+        (opts, Node ns')
+    | Leaf es ->
+        let matching, non_matching = List.partition (fun (_, e) -> eq e) es in
+        let elts = List.map snd matching in
+        (elts, Leaf non_matching)
+    | Empty -> ([], Empty)
+
+  let remove_eq t e =
+    let eq = (Repr.equal V.t |> Repr.unstage) e in
+    match remove_eq' eq t.tree with
+    | [], _ -> None
+    | elts, t' -> Some (elts, { t with tree = t' })
+
+  let rec take_children lst = function
+    | Node ns -> List.split ns |> snd |> List.concat_map (take_children lst)
+    | Leaf es -> List.split es |> snd |> ( @ ) lst
+    | Empty -> []
+
+  let rec remove_env' env = function
+    | Node ns ->
+        let opts, ns' =
+          List.map
+            (fun (e, t) ->
+              if E.contains env e then (take_children [] t, (e, Empty))
+              else
+                let opt, t' = remove_env' env t in
+                (opt, (e, t')))
+            ns
+          |> List.split
+        in
+        let opts = List.concat opts in
+        (opts, Node ns')
+    | Leaf es ->
+        let in_env, out_env =
+          List.partition (fun (e, _) -> E.contains env e) es
+        in
+        (List.map snd in_env, Leaf out_env)
+    | Empty -> ([], Empty)
+
+  let remove_env t env =
+    match remove_env' env t.tree with
+    | [], _ -> None
+    | elts, t' -> Some (elts, { t with tree = t' })
+
   let filter_intersecting e = List.filter (fun (e', _) -> E.intersects e e')
 
   let rec find' t e =

src/rtree_intf.ml

@@ -108,6 +108,14 @@ module type S = sig
   val insert : t -> Value.t -> t
   (** Insert a new element into the tree *)
 
+  val remove_env : t -> Envelope.t -> (Value.t list * t) option
+  (** [remove_env tree envelope] Remove all elements within a certain envelope, 
+      returning them and the new tree or None if none are found *)
+
+  val remove_eq : t -> Value.t -> (Value.t list * t) option
+  (** [remove_eq tree element] Remove all elements with equality to the element passed,
+      returning them and the new tree, or None if none are found *)
+
   val find : t -> Envelope.t -> Value.t list
   (** [find tree env] find all value contained by [env] in [tree]. *)
 

test/basic.ml

@@ -234,6 +234,130 @@ let test_size () =
   let calc_depth = R.size t in
   assert (calc_depth = 4)
 
+let test_remove_eq () =
+  let module R = R1 in
+  let lines =
+    [
+      { p1 = (0., 0.); p2 = (1., 1.) };
+      { p1 = (1., 1.); p2 = (2., 2.) };
+      { p1 = (2., 2.); p2 = (3., 3.) };
+      { p1 = (3., 3.); p2 = (4., 4.) };
+    ]
+  in
+  let idx = R.load ~max_node_load:2 lines in
+  let t' = R.remove_eq idx (List.hd lines) in
+  match t' with
+  | None -> failwith "Unexpected none returned from remove"
+  | Some (vs, t') ->
+      assert (R.size t' = 3);
+      assert (List.hd vs = List.hd lines)
+
+let test_remove_env () =
+  let module R = R1 in
+  let lines =
+    [
+      { p1 = (0., 0.); p2 = (1., 1.) };
+      { p1 = (1., 1.); p2 = (2., 2.) };
+      { p1 = (2., 2.); p2 = (3., 3.) };
+      { p1 = (3., 3.); p2 = (4., 4.) };
+    ]
+  in
+  let idx = R.load ~max_node_load:2 lines in
+  let env = Rtree.Rectangle.v ~x0:0. ~y0:0. ~x1:2. ~y1:2. in
+  let t' = R.remove_env idx env in
+  match t' with
+  | None -> failwith "Unexpected none returned from remove"
+  | Some (vs, t') ->
+      assert (R.size t' = 2);
+      assert (List.length vs = 2)
+
+let test_remove_empty () =
+  let module R = R1 in
+  let empty = R.empty 2 in
+  let env = Rtree.Rectangle.v ~x0:0. ~y0:0. ~x1:1. ~y1:1. in
+  let remove_env = R.remove_env empty env in
+  let line = { p1 = (0., 0.); p2 = (1., 1.) } in
+  let remove_eq = R.remove_eq empty line in
+  assert (Option.is_none remove_env);
+  assert (Option.is_none remove_eq)
+
+let test_remove_one () =
+  let module R = R1 in
+  let line = { p1 = (0., 0.); p2 = (1., 1.) } in
+  let t = R.load [ line ] in
+  let removed_env, t_env =
+    R.remove_env t (Rtree.Rectangle.v ~x0:0. ~y0:0. ~x1:1. ~y1:1.) |> Option.get
+  in
+  let removed_eq, t_eq = R.remove_eq t line |> Option.get in
+  assert (List.length removed_env == 1);
+  assert (R.size t_env == 0);
+  assert (List.length removed_eq == 1);
+  assert (R.size t_eq == 0)
+
+let test_remove_many () =
+  let module R = R1 in
+  let lines =
+    List.init 1_000 (fun _ ->
+        {
+          p1 = (Random.float 1., Random.float 1.);
+          p2 = (Random.float 1., Random.float 1.);
+        })
+  in
+  let t = R.load lines in
+  let env = Rtree.Rectangle.v ~x0:0. ~y0:0. ~x1:1. ~y1:1. in
+  let t_env = R.remove_env t env in
+  let t_eq =
+    List.fold_left
+      (fun acc line ->
+        assert_bool
+          "Remove_many unexpectedly failed: element not present in tree"
+        @@ Option.is_some acc;
+        R.remove_eq (Option.get acc |> snd) line)
+      (Some ([], t))
+      lines
+  in
+  assert (Option.is_some t_env);
+  assert (R.size (Option.get t_env |> snd) == 0);
+  assert (Option.is_some t_eq);
+  assert (R.size (Option.get t_eq |> snd) == 0)
+
+let test_remove_not_present () =
+  let module R = R1 in
+  let lines =
+    List.init 1_000 (fun _ ->
+        {
+          p1 = (Random.float 1., Random.float 1.);
+          p2 = (Random.float 1., Random.float 1.);
+        })
+  in
+  let t = R.load lines in
+  let t_env =
+    R.remove_env t (Rtree.Rectangle.v ~x0:1.01 ~y0:1.01 ~x1:2. ~y1:2.)
+  in
+  let t_eq = R.remove_eq t { p1 = (2., 2.); p2 = (3., 3.) } in
+  assert (Option.is_none t_env);
+  assert (Option.is_none t_eq)
+
+let test_count_env () =
+  let module R = R1 in
+  let ranf () = Random.float 0.5 in
+  let in_range =
+    List.init 50 (fun _ -> { p1 = (ranf (), ranf ()); p2 = (ranf (), ranf ()) })
+  in
+  let out_range =
+    List.init 50 (fun _ ->
+        (* Elts in range starting at .5001 to make sure they're not in the envelope *)
+        {
+          p1 = (0.5001 +. ranf (), 0.5001 +. ranf ());
+          p2 = (0.5001 +. ranf (), 0.5001 +. ranf ());
+        })
+  in
+  let t = R.load (in_range @ out_range) in
+  let env = Rtree.Rectangle.v ~x0:0. ~y0:0. ~x1:0.5 ~y1:0.5 in
+  let t = R.remove_env t env in
+  assert (Option.is_some t);
+  assert (Option.get t |> snd |> R.size |> ( == ) 50)
+
 let suite =
   "R"
   >::: [
@@ -246,6 +370,13 @@ let suite =
          "size" >:: test_size;
          "depth" >:: test_depth;
          "cube" >:: cube;
+         "remove_eq" >:: test_remove_eq;
+         "remove_env" >:: test_remove_env;
+         "remove empty" >:: test_remove_empty;
+         "remove one elt from tree of size one" >:: test_remove_one;
+         "remove many" >:: test_remove_many;
+         "remove element not present in tree" >:: test_remove_not_present;
+         "remove elements within an envelope" >:: test_count_env;
        ]
 
 let _ = run_test_tt_main suite