Skip to content

Commit

Permalink
Fixes the stub resolver tests (#1564)
Browse files Browse the repository at this point in the history 
* Fixes the stub resolver tests

* Preserves transitivity

* Drop `bap-relation`

Co-authored-by: bmourad01 <[email protected]>
  • Loading branch information
@bmourad01
bmourad01 and bmourad01 authored Dec 23, 2022
1 parent ce84c26 commit 79fd1a2
Show file tree
Hide file tree
Showing 3 changed files with 96 additions and 91 deletions.
2 changes: 1 addition & 1 deletion oasis/stub-resolver
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
163 changes: 78 additions & 85 deletions plugins/stub_resolver/stub_resolver.ml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 =
Expand Down Expand Up @@ -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
Expand Down
22 changes: 17 additions & 5 deletions plugins/stub_resolver/stub_resolver_tests.ml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"];
Expand All @@ -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"];
Expand All @@ -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";
];

]

0 comments on commit 79fd1a2

Please sign in to comment.