From 4ac2eb25a67a03dc2b2ff69e8581aa0a980c3a69 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Tue, 7 Dec 2021 21:29:51 -0500
Subject: [PATCH] use TEF in sudoku; improve a bit its ergonomics

---
 examples/sudoku/dune            |  4 +++-
 examples/sudoku/sudoku_solve.ml | 10 ++++++++++
 src/main/main.ml                |  3 +--
 src/tef/Sidekick_tef.dummy.ml   |  1 +
 src/tef/Sidekick_tef.mli        |  2 ++
 src/tef/Sidekick_tef.real.ml    |  4 ++++
 6 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/examples/sudoku/dune b/examples/sudoku/dune
index 8a38ed3e..48d5eb3c 100644
--- a/examples/sudoku/dune
+++ b/examples/sudoku/dune
@@ -2,7 +2,9 @@
 (executable
   (name sudoku_solve)
   (modes native)
-  (libraries sidekick.core sidekick.util sidekick.sat iter containers)
+  (libraries sidekick.core sidekick.util sidekick.sat
+             sidekick.tef ; for profiling
+             iter containers)
   (flags :standard -warn-error -a+8 -color always -safe-string -strict-sequence)
   (ocamlopt_flags :standard -O3 -bin-annot
                   -unbox-closures -unbox-closures-factor 20))
diff --git a/examples/sudoku/sudoku_solve.ml b/examples/sudoku/sudoku_solve.ml
index c6b76f6e..e09522a6 100644
--- a/examples/sudoku/sudoku_solve.ml
+++ b/examples/sudoku/sudoku_solve.ml
@@ -5,6 +5,8 @@ module Fmt = CCFormat
 module Vec = Sidekick_util.Vec
 module Log = Sidekick_util.Log
 
+module Profile = Sidekick_util.Profile
+
 let errorf msg = Fmt.kasprintf failwith msg
 
 module Cell : sig
@@ -173,6 +175,7 @@ end = struct
 
     (* check that all cells are full *)
     let check_full_ (self:t) (acts:(Lit.t,proof,proof_step) acts) : unit =
+      Profile.with_ "check-full" @@ fun () ->
       let (module A) = acts in
       Grid.all_cells (grid self)
         (fun (x,y,c) ->
@@ -187,6 +190,7 @@ end = struct
 
     (* check constraints *)
     let check_ (self:t) (acts:(Lit.t,proof,proof_step) acts) : unit =
+      Profile.with_ "check-constraints" @@ fun () ->
       Log.debugf 4 (fun k->k "(@[sudoku.check@ @[:g %a@]@])" Grid.pp (B_ref.get self.grid));
       let (module A) = acts in
       let[@inline] all_diff kind f =
@@ -237,6 +241,7 @@ end = struct
         )
 
     let partial_check (self:t) acts : unit =
+      Profile.with_ "partial-check" @@ fun () ->
       Log.debugf 4
         (fun k->k "(@[sudoku.partial-check@ :trail [@[%a@]]@])"
             (Fmt.list F.pp) (trail_ acts |> Iter.to_list));
@@ -244,6 +249,7 @@ end = struct
       check_ self acts
 
     let final_check (self:t) acts : unit =
+      Profile.with_ "final-check" @@ fun () ->
       Log.debugf 4 (fun k->k "(@[sudoku.final-check@])");
       check_full_ self acts;
       check_ self acts
@@ -258,6 +264,7 @@ end = struct
   }
 
   let solve (self:t) : _ option =
+    Profile.with_ "sudoku.solve" @@ fun () ->
     let assumptions =
       Grid.all_cells self.grid0
       |> Iter.filter (fun (_,_,c) -> Cell.is_full c)
@@ -296,6 +303,7 @@ let chrono ~pp_time : (module CHRONO) =
   (module M)
 
 let solve_file ~pp_time file =
+  Profile.with_ "solve-file" @@ fun () ->
   let open (val chrono ~pp_time) in
   Format.printf "solve grids in file %S@." file;
 
@@ -331,6 +339,8 @@ let solve_file ~pp_time file =
   ()
 
 let () =
+  Sidekick_tef.with_setup @@ fun () ->
+
   Fmt.set_color_default true;
   let files = ref [] in
   let debug = ref 0 in
diff --git a/src/main/main.ml b/src/main/main.ml
index 14d630e5..d27d8c60 100644
--- a/src/main/main.ml
+++ b/src/main/main.ml
@@ -211,8 +211,7 @@ let main_cnf () : _ result =
 let main () =
 
   (* instrumentation and tracing *)
-  Sidekick_tef.setup();
-  at_exit Sidekick_tef.teardown;
+  Sidekick_tef.with_setup @@ fun () ->
   Sidekick_memtrace.trace_if_requested ~context:"sidekick" ();
 
   CCFormat.set_color_default true;
diff --git a/src/tef/Sidekick_tef.dummy.ml b/src/tef/Sidekick_tef.dummy.ml
index 994edeb7..51c11878 100644
--- a/src/tef/Sidekick_tef.dummy.ml
+++ b/src/tef/Sidekick_tef.dummy.ml
@@ -1,3 +1,4 @@
 
 let setup() = ()
 let teardown() = ()
+let with_setup f = f()
diff --git a/src/tef/Sidekick_tef.mli b/src/tef/Sidekick_tef.mli
index 91cbf736..167a1488 100644
--- a/src/tef/Sidekick_tef.mli
+++ b/src/tef/Sidekick_tef.mli
@@ -22,3 +22,5 @@ val setup : unit -> unit
 (** Install the TEF logger as a profiling backend. *)
 
 val teardown : unit -> unit
+
+val with_setup: (unit -> 'a) -> 'a
diff --git a/src/tef/Sidekick_tef.real.ml b/src/tef/Sidekick_tef.real.ml
index 5b7467b1..a5f5c3a0 100644
--- a/src/tef/Sidekick_tef.real.ml
+++ b/src/tef/Sidekick_tef.real.ml
@@ -72,3 +72,7 @@ let setup_ = lazy (
 
 let setup () = Lazy.force setup_
 let teardown = P.Control.teardown
+let[@inline] with_setup f =
+  setup();
+  try let x = f() in teardown(); x
+  with e -> teardown(); raise e