[Seedpool] refactor construct

Author: bongjunj

src/seedpool.ml

@@ -123,6 +123,7 @@ let check_exist_ret f =
 
 let rec redef_fn llctx f_old wide instr =
   let extra_param = [| ALlvm.i32_type llctx; ALlvm.pointer_type llctx |] in
+  let fmodule = ALlvm.global_parent f_old in
   let params_old = ALlvm.params f_old in
   let param_tys =
     Array.append (Array.map ALlvm.type_of params_old) extra_param
@@ -134,11 +135,8 @@ let rec redef_fn llctx f_old wide instr =
       if i = instr then true
       else if ALlvm.ChangeRetVal.check_target i then (
         let new_ret_ty = ALlvm.type_of i in
-        let f_new =
-          ALlvm.define_function ""
-            (ALlvm.function_type new_ret_ty param_tys)
-            (ALlvm.global_parent f_old)
-        in
+        let new_fty = ALlvm.function_type new_ret_ty param_tys in
+        let f_new = ALlvm.define_function "" new_fty fmodule in
 
         (* copy function with new return value (target).*)
         ALlvm.ChangeRetVal.copy_blocks llctx f_old f_new;
@@ -293,90 +291,83 @@ let save ?(parent : int option) (seed : Seed.t) =
       Printf.fprintf line "%s" seed.importants);
   ALlvm.save_ll !Config.corpus_dir seed_name llm |> ignore
 
-let evaluate_seeds_and_construct_seedpool seeds node_tbl distmap =
+let is_optimuzz_base () =
+  (* Optimuzz_Base *)
+  !Config.score = Config.FuzzingMode.Constant
+  && !Config.coverage = Config.FuzzingMode.All_edges
+
+let is_nondirected () = !Config.coverage = Config.FuzzingMode.All_edges
+
+let log_seed seed =
+  L.debug "evaluate seed: \n%s\n" (ALlvm.string_of_llmodule seed.Seed.llm);
+  L.debug "score: %f" (Seed.score seed)
+
+let classify_seeds node_tbl distmap raw_seeds =
+  if is_nondirected () then
+    List.fold_left
+      (fun (cover_seeds, noncover_seeds) (_, llm, traces) ->
+        let seed = Seed.make llm traces "" node_tbl distmap in
+        log_seed seed;
+        (cover_seeds, seed :: noncover_seeds))
+      ([], []) raw_seeds
+  else
+    List.fold_left
+      (fun (cover_seeds, noncover_seeds) (_, llm, traces) ->
+        let seed = Seed.make llm traces "" node_tbl distmap in
+        log_seed seed;
+        if Seed.covers seed then (seed :: cover_seeds, noncover_seeds)
+        else (cover_seeds, seed :: noncover_seeds))
+      ([], []) raw_seeds
+
+let hash_seed seed = ALlvm.hash_llm seed.Seed.llm
+let compare_hash a b = compare (hash_seed a) (hash_seed b)
+let compare_score a b = compare a.Seed.score b.Seed.score
+
+let closest_seeds seeds =
+  let unique_seeds = List.sort_uniq compare_hash seeds in
+  let sorted_seeds =
+    if !Config.score = Config.FuzzingMode.Constant then unique_seeds
+    else
+      unique_seeds
+      |> List.sort (fun a b -> compare (Seed.score a) (Seed.score b))
+  in
+  let first_seeds =
+    if !Config.coverage = Config.FuzzingMode.All_edges then sorted_seeds
+    else List.take !Config.max_initial_seed sorted_seeds
+  in
+  first_seeds
+
+let compute_init_cov seeds =
+  let open Coverage in
+  List.fold_left
+    (fun accu seed -> EdgeCoverage.union accu seed.Seed.edge_cov)
+    EdgeCoverage.empty seeds
+
+let evaluate_seeds_and_construct_seedpool raw_seeds node_tbl distmap =
   let open AUtil in
+  let open Coverage in
   let pool = create () in
-  if
-    (* Optimuzz_Base *)
-    !Config.score = Config.FuzzingMode.Constant
-    && !Config.coverage = Config.FuzzingMode.All_edges
-  then (
-    let pools =
-      List.fold_left
-        (fun pools (_, llm, traces) ->
-          let seed = Seed.make llm traces "" node_tbl distmap in
-          L.debug "evaluate seed: \n%s\n" (ALlvm.string_of_llmodule llm);
-          L.debug "score: %s" (string_of_float (Seed.score seed));
-          seed :: pools)
-        [] seeds
-    in
-    let init_cov =
-      List.fold_left
-        (fun accu seed -> Coverage.EdgeCoverage.union accu (Seed.edge_cov seed))
-        Coverage.EdgeCoverage.empty pools
-    in
-    pools |> List.to_seq |> add_seq pool;
+  if is_optimuzz_base () then (
+    let _, seeds = classify_seeds node_tbl distmap raw_seeds in
+    let init_cov = compute_init_cov seeds in
+    seeds |> List.to_seq |> add_seq pool;
     (pool, init_cov))
   else
-    let pool_covers, pool_noncovers =
-      List.fold_left
-        (fun (pool_covers, pool_noncovers) (_, llm, traces) ->
-          let seed = Seed.make llm traces "" node_tbl distmap in
-          L.debug "evaluate seed: \n%s\n" (ALlvm.string_of_llmodule llm);
-          L.debug "score: %s" (string_of_float (Seed.score seed));
-          (* will assume all-edges option to non-directed fuzzing *)
-          if !Config.coverage = Config.FuzzingMode.All_edges then
-            (pool_covers, seed :: pool_noncovers)
-          else if Seed.covers seed then (seed :: pool_covers, pool_noncovers)
-          else (pool_covers, seed :: pool_noncovers))
-        ([], []) seeds
+    let cover_seeds, noncover_seeds =
+      classify_seeds node_tbl distmap raw_seeds
     in
-    if pool_covers = [] then (
+    if cover_seeds = [] then (
       L.info "No covering seeds found. Using closest seeds.";
-      let pool_closest =
-        pool_noncovers
-        |> List.sort_uniq (fun a b ->
-               compare
-                 (ALlvm.hash_llm (Seed.llmodule a))
-                 (ALlvm.hash_llm (Seed.llmodule b)))
-      in
-      let pool_closest =
-        if !Config.score = Config.FuzzingMode.Constant then pool_closest
-        else
-          pool_closest
-          |> List.sort (fun a b -> compare (Seed.score a) (Seed.score b))
-      in
-      let pool_closest =
-        (* will assume all-edges option to non-directed fuzzing *)
-        if !Config.coverage = Config.FuzzingMode.All_edges then pool_closest
-        else pool_closest |> take !Config.max_initial_seed
-      in
-
-      let init_cov =
-        List.fold_left
-          (fun accu seed ->
-            Coverage.EdgeCoverage.union accu (Seed.edge_cov seed))
-          Coverage.EdgeCoverage.empty pool_closest
-      in
+      let pool_closest = closest_seeds noncover_seeds in
+      let init_cov = compute_init_cov pool_closest in
       pool_closest |> List.to_seq |> add_seq pool;
       (pool, init_cov))
     else (
       (* if we have covering seeds, we use covering seeds only. *)
       L.info "Covering seeds found. Using them only.";
-      let pool_covers =
-        pool_covers
-        |> List.sort_uniq (fun a b ->
-               compare
-                 (ALlvm.hash_llm (Seed.llmodule a))
-                 (ALlvm.hash_llm (Seed.llmodule b)))
-      in
-      let init_cov =
-        List.fold_left
-          (fun accu seed ->
-            Coverage.EdgeCoverage.union accu (Seed.edge_cov seed))
-          Coverage.EdgeCoverage.empty pool_covers
-      in
-      pool_covers |> List.to_seq |> add_seq pool;
+      let sorted_seeds = cover_seeds |> List.sort_uniq compare_score in
+      let init_cov = compute_init_cov sorted_seeds in
+      sorted_seeds |> List.to_seq |> add_seq pool;
       (pool, init_cov))
 
 let make llctx node_tbl (distmap : float Coverage.DistanceTable.t) =
@@ -412,7 +403,7 @@ let make llctx node_tbl (distmap : float Coverage.DistanceTable.t) =
 
   let seed_count = ref 0 in
 
-  let filter_seed seed =
+  let preprocess_seed seed =
     let* path, lines = can_optimize seed node_tbl distmap in
     L.debug "filter seed: %s " path;
     match ALlvm.read_ll llctx path with
@@ -432,7 +423,7 @@ let make llctx node_tbl (distmap : float Coverage.DistanceTable.t) =
 
     Array.to_list seed_files
     |> List.map (Filename.concat seed_dir)
-    |> List.filter_map filter_seed
+    |> List.filter_map preprocess_seed
   in
 
   let pool = evaluate_seeds_and_construct_seedpool seeds node_tbl distmap in