[Fuzzer] refactor

Author: bongjunj

src/coverage.ml

@@ -56,11 +56,7 @@ end
 module EdgeCoverage = struct
   include Set.Make (IntInt)
 
-  let of_traces traces =
-    traces
-    |> List.map AUtil.pairs
-    |> List.map of_list
-    |> List.fold_left union empty
+  let of_trace (trace : int list) = of_list (AUtil.pairs trace)
 
   let read file =
     let open AUtil in
@@ -439,7 +435,7 @@ module DistanceTable = struct
     |> map harmonic_avg
 end
 
-let sliced_cfg_node_of_addr node_tbl distmap addr =
+let node_of_addr node_tbl distmap addr =
   match AddrToNode.find_opt node_tbl addr with
   | None -> None (* in CFG of a function other than the target function *)
   | Some node -> if DistanceTable.mem node distmap then Some node else None
@@ -455,37 +451,23 @@ module CfgDistance = struct
   type t = float
 
   (** computes distance of a trace *)
-  let distance_score (traces : BlockTrace.t list) node_tbl distmap =
-    let nodes_in_trace =
-      traces
-      |> List.flatten
-      |> List.sort_uniq compare
-      |> List.filter_map (sliced_cfg_node_of_addr node_tbl distmap)
-    in
-    let dist_sum : float =
-      nodes_in_trace
-      |> List.fold_left
-           (fun sum node ->
-             let dist = DistanceTable.find_opt node distmap in
-             match dist with None -> sum | Some dist -> sum +. dist)
-           0.0
-    in
-    (* let min_dist =
-           nodes_in_trace
-           |> List.filter_map (fun node -> Cfg.NodeMap.find_opt node distmap)
-           |> List.fold_left (fun accu dist -> Float.min accu dist) 65535.0
-         in *)
-    if nodes_in_trace = [] then 65535.0
+  let distance_score (trace : BlockTrace.t) node_tbl distmap =
+    let nodes = trace |> List.filter_map (node_of_addr node_tbl distmap) in
+    if nodes = [] then 65535.0
     else
-      let size = List.length nodes_in_trace |> float_of_int in
-      dist_sum /. size
-
-  let get_cover (traces : BlockTrace.t list) node_tbl distmap =
-    traces
-    |> List.exists
-         (List.exists (fun addr ->
-              let node = AddrToNode.find_opt node_tbl addr in
-              match node with
-              | None -> false
-              | Some node -> DistanceTable.find_opt node distmap = Some 0.0))
+      let dist_sum : float =
+        nodes
+        |> List.filter_map (fun node -> DistanceTable.find_opt node distmap)
+        |> List.fold_left (fun sum dist -> sum +. dist) 0.0
+      in
+      let size = List.length nodes in
+      dist_sum /. float_of_int size
+
+  let get_cover (trace : BlockTrace.t) node_tbl distmap =
+    trace
+    |> List.exists (fun addr ->
+           let node = AddrToNode.find_opt node_tbl addr in
+           match node with
+           | None -> false
+           | Some node -> DistanceTable.find_opt node distmap = Some 0.0)
 end

src/fuzzer.ml

@@ -31,35 +31,37 @@ let measure_optimizer_coverage llm =
     AUtil.clean optimized_ir_filename;
     (optimization_res, validation_res)
 
-let evaluate_mutant parent_llm llm importants covset node_tbl distance_map =
+let is_in_slice addr node_tbl distmap =
+  if !Config.coverage = Config.FuzzingMode.Sliced_cfg then
+    Coverage.node_of_addr node_tbl distmap addr |> Option.is_some
+  else true
+
+let log_new_points new_points =
+  if Coverage.EdgeCoverage.is_empty new_points then
+    L.debug "No new coverage points"
+  else L.debug "New coverage points"
+
+let save_covering_mutant parent_llm seed =
+  if seed.Seedpool.CfgSeed.covers then
+    let seed_name =
+      Seedpool.Seed.name ~parent:(ALlvm.hash_llm parent_llm) seed
+    in
+    ALlvm.save_ll !Config.covers_dir seed_name seed.llm |> ignore
+
+let evaluate_mutant parent_llm llm importants covset node_tbl distmap =
   let optim_res, _ = measure_optimizer_coverage llm in
   L.debug "Mutant: ";
   L.debug "%s" (ALlvm.string_of_llmodule llm);
   match optim_res with
   | Error _ -> None
   | Ok lines ->
-      let in_slice =
-        match !Config.coverage with
-        | Config.FuzzingMode.Sliced_cfg ->
-            fun addr ->
-              Coverage.sliced_cfg_node_of_addr node_tbl distance_map addr
-              |> Option.is_some
-        | _ -> fun _ -> true
-      in
-      let trace = lines |> List.map int_of_string |> List.filter in_slice in
+      let is_in_slice addr = is_in_slice addr node_tbl distmap in
+      let trace = lines |> List.map int_of_string |> List.filter is_in_slice in
       let cov = trace |> AUtil.pairs |> Coverage.EdgeCoverage.of_list in
       let new_points = Coverage.EdgeCoverage.diff cov covset in
-      if Coverage.EdgeCoverage.is_empty new_points then
-        prerr_endline "No new coverage points"
-      else prerr_endline "New coverage points";
-      let (new_seed : Seedpool.Seed.t) =
-        Seedpool.Seed.make llm [ trace ] importants node_tbl distance_map
-      in
-      (if new_seed.covers then
-         let seed_name =
-           Seedpool.Seed.name ~parent:(ALlvm.hash_llm parent_llm) new_seed
-         in
-         ALlvm.save_ll !Config.covers_dir seed_name llm |> ignore);
+      log_new_points new_points;
+      let new_seed = Seedpool.Seed.make llm trace importants node_tbl distmap in
+      save_covering_mutant parent_llm new_seed;
       let interesting = not (Coverage.EdgeCoverage.is_empty new_points) in
       if interesting then Some new_seed else None
   | Assert _ -> None
@@ -97,13 +99,14 @@ let update_progress progress seed =
   let cov_set = Seedpool.Seed.edge_cov seed in
   progress |> Progress.add_cov cov_set |> Progress.inc_gen
 
-let rec try_gen_mutant mutator node_tbl distmap energy llm (prog : Progress.t) =
+let rec try_gen_mutant mutator node_tbl distmap energy llm prog =
   if energy = 0 then None
   else
     match mutator llm with
     | Some (mutant, importants) -> (
         match
-          evaluate_mutant llm mutant importants prog.cov_sofar node_tbl distmap
+          evaluate_mutant llm mutant importants prog.Progress.cov_sofar node_tbl
+            distmap
         with
         | Some new_seed -> Some new_seed
         | None ->

src/seedpool.ml

@@ -32,15 +32,14 @@ module CfgSeed = struct
     importants : string;
   }
 
-  let make llm (traces : Coverage.BlockTrace.t list) importants node_tbl distmap
-      =
-    let score = Distance.distance_score traces node_tbl distmap in
-    let covers = Distance.get_cover traces node_tbl distmap in
+  let make llm (trace : Coverage.BlockTrace.t) importants node_tbl distmap =
+    let score = Distance.distance_score trace node_tbl distmap in
+    let covers = Distance.get_cover trace node_tbl distmap in
     {
       llm;
       score;
       covers;
-      edge_cov = Coverage.EdgeCoverage.of_traces traces;
+      edge_cov = Coverage.EdgeCoverage.of_trace trace;
       importants;
     }
 
@@ -285,7 +284,7 @@ let can_optimize seedfile node_tbl distmap =
         | Config.FuzzingMode.Sliced_cfg ->
             fun line ->
               int_of_string line
-              |> Coverage.sliced_cfg_node_of_addr node_tbl distmap
+              |> Coverage.node_of_addr node_tbl distmap
               |> Option.is_some
         | _ -> fun _ -> true
       in
@@ -441,7 +440,7 @@ let make llctx node_tbl (distmap : float Coverage.DistanceTable.t) =
     | Error _ -> None
     | Ok llm ->
         if check_llm_for_mutation llm then (
-          let cov = [ lines |> List.map int_of_string ] in
+          let cov = lines |> List.map int_of_string in
           let llm = add_dummy_params llm in
           L.debug "filtered seeds: %s" (ALlvm.string_of_llmodule llm);
           seed_count := !seed_count + 1;