From d9baa2f542dd6fd0d60d12dfdf3ba1df5032b12b Mon Sep 17 00:00:00 2001
From: Simmo Saan <simmo.saan@gmail.com>
Date: Mon, 6 Nov 2023 13:46:13 +0200
Subject: [PATCH] Add WidenedVars query

---
 src/analyses/base.ml       |  2 ++
 src/cdomains/baseDomain.ml |  5 +++--
 src/domains/queries.ml     |  5 +++++
 src/witness/yamlWitness.ml | 13 +++++++++----
 4 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/analyses/base.ml b/src/analyses/base.ml
index a8ad9af95b..3b727edc06 100644
--- a/src/analyses/base.ml
+++ b/src/analyses/base.ml
@@ -1378,6 +1378,8 @@ struct
     | Q.InvariantGlobal g ->
       let g: V.t = Obj.obj g in
       query_invariant_global ctx g
+    | Q.WidenedVars ->
+      snd ctx.local.cpa
     | _ -> Q.Result.top q
 
   let update_variable variable typ value cpa =
diff --git a/src/cdomains/baseDomain.ml b/src/cdomains/baseDomain.ml
index 0372e38165..379fbe569d 100644
--- a/src/cdomains/baseDomain.ml
+++ b/src/cdomains/baseDomain.ml
@@ -19,10 +19,11 @@ struct
   let name () = "value domain"
 end
 
-module CPA: MapDomain.S with type key = Basetype.Variables.t and type value = VD.t =
+module Vars = Queries.VS
+
+module CPA (* : MapDomain.S with type key = Basetype.Variables.t and type value = VD.t *) =
 struct
   module CPA = CPA0
-  module Vars = SetDomain.ToppedSet (Basetype.Variables) (struct let topname = "Top" end)
   include Lattice.Prod (CPA) (Vars)
 
   type key = CPA.key
diff --git a/src/domains/queries.ml b/src/domains/queries.ml
index c706339bf2..41baf874c3 100644
--- a/src/domains/queries.ml
+++ b/src/domains/queries.ml
@@ -128,6 +128,7 @@ type _ t =
   | MayBeTainted: AD.t t
   | MayBeModifiedSinceSetjmp: JmpBufDomain.BufferEntry.t -> VS.t t
   | TmpSpecial:  Mval.Exp.t -> ML.t t
+  | WidenedVars: VS.t t
 
 type 'a result = 'a
 
@@ -194,6 +195,7 @@ struct
     | MayBeTainted -> (module AD)
     | MayBeModifiedSinceSetjmp _ -> (module VS)
     | TmpSpecial _ -> (module ML)
+    | WidenedVars -> (module VS)
 
   (** Get bottom result for query. *)
   let bot (type a) (q: a t): a result =
@@ -259,6 +261,7 @@ struct
     | MayBeTainted -> AD.top ()
     | MayBeModifiedSinceSetjmp _ -> VS.top ()
     | TmpSpecial _ -> ML.top ()
+    | WidenedVars -> VS.top ()
 end
 
 (* The type any_query can't be directly defined in Any as t,
@@ -321,6 +324,7 @@ struct
     | Any ThreadsJoinedCleanly -> 52
     | Any (TmpSpecial _) -> 53
     | Any (IsAllocVar _) -> 54
+    | Any WidenedVars -> 55
 
   let rec compare a b =
     let r = Stdlib.compare (order a) (order b) in
@@ -472,6 +476,7 @@ struct
     | Any DYojson -> Pretty.dprintf "DYojson"
     | Any MayBeModifiedSinceSetjmp buf -> Pretty.dprintf "MayBeModifiedSinceSetjmp %a" JmpBufDomain.BufferEntry.pretty buf
     | Any (TmpSpecial lv) -> Pretty.dprintf "TmpSpecial %a" Mval.Exp.pretty lv
+    | Any WidenedVars -> Pretty.dprintf "WidenedVars"
 end
 
 let to_value_domain_ask (ask: ask) =
diff --git a/src/witness/yamlWitness.ml b/src/witness/yamlWitness.ml
index ef919942c7..4aed973a36 100644
--- a/src/witness/yamlWitness.ml
+++ b/src/witness/yamlWitness.ml
@@ -200,10 +200,15 @@ struct
           lvals
       )
       else (
-        let vars = BaseDomain.NH.find_all BaseDomain.widen_vars n in
-        List.fold_left (fun acc var ->
-            Lval.Set.add (Var var, NoOffset) acc
-          ) (Lval.Set.empty ()) vars
+        let vars = R.ask_local_node n ~local WidenedVars in
+        ignore (Pretty.printf "yaml widened vars %a at %a\n" Queries.VS.pretty vars Node.pretty_plain_short n);
+        if Queries.VS.is_top vars then
+          Lval.Set.top ()
+        else (
+          Queries.VS.fold (fun var acc ->
+              Lval.Set.add (Var var, NoOffset) acc
+            ) vars (Lval.Set.empty ())
+        )
       )
     in