From 79fd1a2fe74e87b7be5da35accdac7c95df46d21 Mon Sep 17 00:00:00 2001 From: Benjamin Mourad Date: Fri, 23 Dec 2022 10:57:24 -0500 Subject: [PATCH] Fixes the stub resolver tests (#1564) * Fixes the stub resolver tests * Preserves transitivity * Drop `bap-relation` Co-authored-by: bmourad01 --- oasis/stub-resolver | 2 +- plugins/stub_resolver/stub_resolver.ml | 163 +++++++++---------- plugins/stub_resolver/stub_resolver_tests.ml | 22 ++- 3 files changed, 96 insertions(+), 91 deletions(-) diff --git a/oasis/stub-resolver b/oasis/stub-resolver index fa452f554..3564cfa20 100644 --- a/oasis/stub-resolver +++ b/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 diff --git a/plugins/stub_resolver/stub_resolver.ml b/plugins/stub_resolver/stub_resolver.ml index 767a5bde2..1b50f8751 100644 --- a/plugins/stub_resolver/stub_resolver.ml +++ b/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 diff --git a/plugins/stub_resolver/stub_resolver_tests.ml b/plugins/stub_resolver/stub_resolver_tests.ml index 5e7112298..85357f065 100644 --- a/plugins/stub_resolver/stub_resolver_tests.ml +++ b/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,7 +228,7 @@ 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"]; @@ -228,6 +236,10 @@ let suite = "stub-resolver" >::: [ real "p5" []; stub "p6" ["p8"; "p9"; "p10"; "p4"]; real "p11" ["p12"; "p13"; "p1"]; - ] ~expected:["p0", "p11" ]; + ] ~expected:[ + "p0", "p11"; + "p4", "p5"; + "p6", "p5"; + ]; ]