Fixes the stub resolver tests (#1564)

* Fixes the stub resolver tests

* Preserves transitivity

* Drop `bap-relation`

Co-authored-by: bmourad01 <[email protected]>

Author: bmourad01

oasis/stub-resolver

@@ -7,7 +7,7 @@ Library bap-plugin-stub_resolver
   Path:         plugins/stub_resolver/
   BuildDepends: bap, bap-abi, bap-knowledge, bap-core-theory, core_kernel,
                 bitvec, bitvec-order, bitvec-sexp,
-                bap-main, ppx_bap, ogre
+                bap-main, ppx_bap, ogre, graphlib, regular
   FindlibName:     bap-plugin-stub_resolver
   CompiledObject:  best
   Modules:         Stub_resolver

plugins/stub_resolver/stub_resolver.ml

@@ -2,21 +2,30 @@ open Core_kernel[@@warning "-D"]
 open Bap.Std
 open Bap_core_theory
 open Bap_knowledge
+open Graphlib.Std
+open Regular.Std
 
 include Self ()
 let package = "bap"
 
 open KB.Syntax
 
-type groups = int Tid.Map.t
-type names = String.Set.t Int.Map.t
+module Regular_string = struct
+  type t = string
+  include Regular.Make(struct
+      include String
+      let module_name = Some "String"
+      let version = "2.6.0"
+    end)
+end
+
+module G = Graphlib.Make(Regular_string)(Unit)
 
 type state = {
-  groups : groups;
-  names  : names;
-  next   : int;
-  stubs  : Tid.Set.t;
-  units  : Theory.Unit.t Tid.Map.t;
+  graph : G.t;
+  names : Tid.Set.t String.Map.t;
+  stubs : Tid.Set.t;
+  units : Theory.Unit.t Tid.Map.t;
 }
 
 module Class = struct
@@ -39,11 +48,10 @@ module Class = struct
 end
 
 let empty = {
-  groups = Map.empty (module Tid);
-  names  = Map.empty (module Int);
-  units  = Map.empty (module Tid);
-  stubs  = Set.empty (module Tid);
-  next   = 0;
+  graph = G.empty;
+  names = String.Map.empty;
+  stubs = Tid.Set.empty;
+  units = Tid.Map.empty;
 }
 
 let in_file file f =
@@ -83,84 +91,69 @@ let update_units t sub =
     !!{t with units = Map.add_exn t.units tid unit}
   | None -> !!t
 
-let find_groups names aliases =
-  Map.fold names ~init:[]
-    ~f:(fun ~key:group ~data:aliases' groups ->
-        if Set.(is_empty @@ inter aliases aliases')
-        then groups
-        else group :: groups)
-
-let unite_names t groups =
-  List.fold groups ~init:(Set.empty (module String))
-    ~f:(fun als id ->
-        Set.union als (Map.find_exn t.names id))
-
-let pick_representative = function
-  | [] -> assert false
-  | groups ->
-    Option.value_exn (List.min_elt groups ~compare:Int.compare)
-
-
-let redirect t ~from ~to_ =
-  Map.map t.groups ~f:(fun id ->
-      if List.mem from id ~equal:Int.equal
-      then to_
-      else id)
-
-let add t sub =
+let should_link aliases ~link_only ~no_link =
+  Set.(is_empty @@ inter aliases no_link) && begin
+    Set.is_empty link_only ||
+    not Set.(is_empty @@ inter aliases link_only)
+  end
+
+let update_graph t name aliases =
+  let n = G.Node.create name in
+  let init = G.Node.insert n t.graph in
+  let graph = Set.fold aliases ~init ~f:(fun g alias ->
+      if String.(name <> alias) then
+        let a = G.Node.create alias in
+        let x = G.Edge.create n a () in
+        let y = G.Edge.create a n () in
+        G.Edge.(insert x (insert y g))
+      else g) in
+  {t with graph}
+
+let update_names t sub ~link_only ~no_link =
+  aliases_of_sub sub >>| fun aliases ->
+  if should_link aliases ~link_only ~no_link then
+    let tid = Term.tid sub in
+    let names = Set.fold aliases ~init:t.names ~f:(fun m a ->
+        Map.update m a ~f:(function
+            | None -> Tid.Set.singleton tid
+            | Some s -> Set.add s tid)) in
+    update_graph {t with names} (Sub.name sub) aliases
+  else t
+
+let add t sub ~link_only ~no_link =
   update_stubs t sub >>= fun t ->
   update_units t sub >>= fun t ->
-  aliases_of_sub sub >>| fun aliases ->
-  match find_groups t.names aliases with
-  | [] ->
-    let groups = Map.add_exn t.groups (Term.tid sub) t.next in
-    let names  = Map.add_exn t.names t.next aliases in
-    {t with groups; names; next = t.next + 1}
-  | [id] ->
-    let groups = Map.add_exn t.groups (Term.tid sub) id in
-    let names = Map.update t.names id ~f:(function
-        | None -> assert false
-        | Some als' -> Set.union aliases als') in
-    {t with names; groups}
-  | groups ->
-    let grp = pick_representative groups in
-    let aliases = Set.union aliases (unite_names t groups) in
-    let names = List.fold groups ~init:t.names ~f:Map.remove in
-    let names = Map.add_exn names ~key:grp ~data:aliases in
-    let groups = redirect t ~from:groups ~to_:grp in
-    {t with names; groups}
-
-let collect_by_group_id stubs groups =
-  Map.fold groups ~init:Int.Map.empty
-    ~f:(fun ~key:tid ~data:id xs ->
-        Map.update xs id ~f:(function
-            | None -> [tid]
-            | Some tids -> tid :: tids)) |>
-  Map.map ~f:(List.partition_tf ~f:(Set.mem stubs))
-
-let unambiguous_pairs names xs ~link_only ~no_link =
-  let should_link id names =
-    let names = Map.find_exn names id in
-    Set.(is_empty @@ inter names no_link) && begin
-      Set.is_empty link_only ||
-      not Set.(is_empty @@ inter names link_only)
-    end in
-  let add y pairs x = Map.add_exn pairs x y in
-  Map.fold xs ~init:(Map.empty (module Tid))
-    ~f:(fun ~key:id ~data:(stubs, impls) init ->
-        match impls with
-        | [y] when should_link id names ->
-          List.fold stubs ~init ~f:(add y)
-        | _ -> init)
-
-let find_pairs t ~link_only ~no_link =
-  unambiguous_pairs t.names ~link_only ~no_link @@
-  collect_by_group_id t.stubs t.groups
+  update_names t sub ~link_only ~no_link
+
+let partition_group t group =
+  Group.enum group |>
+  Seq.fold ~init:Tid.Set.empty ~f:(fun default name ->
+      Map.find t.names name |>
+      Option.value_map ~default ~f:(Set.union default)) |>
+  Set.partition_tf ~f:(Set.mem t.stubs)
+
+let find_pairs t =
+  let pp = Group.pp String.pp in
+  Graphlib.strong_components (module G) t.graph |>
+  Partition.groups |> Seq.fold ~init:Tid.Map.empty ~f:(fun init group ->
+      let stubs, reals = partition_group t group in
+      match Set.length reals with
+      | 1 ->
+        let impl = Set.min_elt_exn reals in
+        Set.fold stubs ~init ~f:(fun links stub ->
+            Map.add_exn links stub impl)
+      | 0 ->
+        info "no implementations found in group %a" pp group;
+        init
+      | n ->
+        info "ambiguous implementations (%d) found in group %a" n pp group;
+        init)
 
 let resolve prog ~link_only ~no_link =
+  let f = add ~link_only ~no_link in
   Term.to_sequence sub_t prog |>
-  Knowledge.Seq.fold ~init:empty ~f:add >>| fun state ->
-  state, find_pairs state ~link_only ~no_link
+  Knowledge.Seq.fold ~init:empty ~f >>| fun state ->
+  state, find_pairs state
 
 let label_name x =
   KB.collect Theory.Label.name x >>| function

plugins/stub_resolver/stub_resolver_tests.ml

@@ -177,13 +177,19 @@ let suite = "stub-resolver" >::: [
       real "h0" ["h1"; "h2"];
       stub "h1" [];
       stub "h2" [];
-    ] ~expected:[];
+    ] ~expected:[
+      "h1", "h0";
+      "h2", "h0";
+    ];
 
     test "ambiguous stubs" [
       real "i0" [];
       stub "i1" ["i0"];
       stub "i2" ["i0"];
-    ] ~expected:[];
+    ] ~expected:[
+      "i1", "i0";
+      "i2", "i0";
+    ];
 
     test "crossreference" [
       real "j0" ["j1"];
@@ -210,7 +216,9 @@ let suite = "stub-resolver" >::: [
       real "m6" ["m5"; "m9"];
       real "m9" ["m10"];
       stub "m10" [];
-    ] ~expected:["m0", "m1"; ];
+    ] ~expected:[
+      "m0", "m1";
+    ];
 
     test "several intersections 2" [
       stub "n0" ["n1"; "n2"; "n3"];
@@ -220,14 +228,18 @@ let suite = "stub-resolver" >::: [
       real "n6" [];
       stub "n7" ["n6"];
       real "n8" ["n1"; "n5"]
-    ] ~expected:["n7", "n6" ];
+    ] ~expected:["n7", "n6"];
 
     test "several intersections 3" [
       stub "p0" ["p1"; "p2"; "p3"];
       stub "p4" ["p5"; "p6"; "p7"];
       real "p5" [];
       stub "p6" ["p8"; "p9"; "p10"; "p4"];
       real "p11" ["p12"; "p13"; "p1"];
-    ] ~expected:["p0", "p11" ];
+    ] ~expected:[
+      "p0", "p11";
+      "p4", "p5";
+      "p6", "p5";
+    ];
 
   ]