Add RecordAscriptionProbe effect and integrate sample recording in transitions

Author: 7h3kk1d

src/language/dynamics/state/EvaluatorState.re

@@ -17,6 +17,7 @@ type effect =
   | RecordStackFrame
   | RecordPatProbes(PatternMatch.sample_closures)
   | RecordTypeInstantiation(Sample.call_stack => Dynamics.TypeInstantiation.t)
+  | RecordAscriptionProbe((Id.t, Sample.capture_spec, Exp.t))
   | RecordTheorem(Id.t, string, Environment.t(Exp.t), Exp.t)
   | RecordPrint(DHExp.t); /* Println for probes study */
 
@@ -145,6 +146,27 @@ let update =
           step_count: state.step_count + 1,
         };
         (call_stack, state);
+      | RecordAscriptionProbe((id, capture_spec, ascribed_exp)) =>
+        let step = state.step_count;
+        let sample =
+          Sample.mk(
+            ~step_start=step,
+            ~step_end=step,
+            id,
+            ascribed_exp,
+            env,
+            call_stack,
+            capture_spec,
+          );
+        print_endline("Adding sample");
+        let state = add_sample(state, sample);
+        /* Advance step count past ascription evaluation */
+        // TODO Check with disconcision
+        let state = {
+          ...state,
+          step_count: state.step_count + 1,
+        };
+        (call_stack, state);
       | RecordTypeInstantiation(type_inst_closure) => (
           call_stack,
           add_type_inst(state, type_inst_closure(call_stack)),

src/language/dynamics/transition/Ascriptions.re

@@ -18,6 +18,7 @@
  when constructing a new expression structure.
  */
 // TODO Rename to samples
+[@deriving show]
 type ascription_samples = list((Id.t, Sample.capture_spec, Exp.t));
 module SampleWriter =
   Util.WriterMonad.Make({
@@ -40,15 +41,20 @@ let rec transition =
   switch (DHExp.term_of(d)) {
   | Asc(e, t) =>
     switch (DHExp.term_of(e), Typ.term_of(Typ.unroll(t))) {
-    | (Asc(e, t'), _)
+    | (Asc(e', t'), _)
         // This is only necessary because sometimes we add two ascriptions and aren't marking it as a non-value
         when Typ.is_consistent(Ctx.empty, Typ.unroll(t), Typ.unroll(t')) =>
+      let* () =
+        switch (Id.Map.find_opt(Exp.rep_id(e), targets)) {
+        | Some(spec) => SampleWriter.tell([(Exp.rep_id(e), spec, e')])
+        | None => SampleWriter.return()
+        };
       switch (Typ.meet(Ctx.empty, Typ.unroll(t), Typ.unroll(t'))) {
       | Some(t) =>
-        let+ result = recur(Asc(e, t) |> DHExp.fresh);
+        let+ result = recur(Asc(e', t) |> DHExp.fresh);
         Some(result);
       | None => SampleWriter.return(None) //TODO  This is an impossible case since we checked consistency
-      }
+      };
     | (e, Parens(t)) =>
       // This is an impossible case since types should be normalized before coming to transitions
       transition(
@@ -298,8 +304,6 @@ let rec transition_multiple =
   };
 };
 
-let transition = (~targets: Sample.targets, d: DHExp.t): option(DHExp.t) =>
-  snd(transition(~targets, ~recursive=false, d));
-
+// TODO return samples
 let transition_multiple = (~targets: Sample.targets, d: DHExp.t): DHExp.t =>
   snd(transition_multiple(~targets, d));

src/language/dynamics/transition/Transition.re

@@ -937,29 +937,37 @@ module Transition = (EV: EV_MODE) => {
       let.wrap_closure _ = (env, d);
       Indet;
     | Asc(d', t) =>
-      switch (Ascriptions.transition(~targets, d)) {
+      let (samples, stepped) = Ascriptions.transition(~targets, d);
+      switch (stepped) {
       | Some(d') =>
         let. _ = otherwise(env, d);
         Step({
           expr: d',
-          side_effects: [],
+          side_effects:
+            List.map(x => EvaluatorState.RecordAscriptionProbe(x), samples),
           kind: Ascription,
           is_value: false,
         });
       | None =>
         let. _ = otherwise(env, d => Asc(d, t) |> rewrap)
         and. d' = req_final(req(env), d => Asc(d, t) |> wrap_ctx, d');
-        switch (Ascriptions.transition(~targets, Asc(d', t) |> rewrap)) {
+        let (samples, stepped) =
+          Ascriptions.transition(~targets, Asc(d', t) |> rewrap);
+        switch (stepped) {
         | Some(d) =>
           Step({
             expr: d,
-            side_effects: [],
+            side_effects:
+              List.map(
+                x => EvaluatorState.RecordAscriptionProbe(x),
+                samples,
+              ),
             kind: Ascription,
             is_value: false,
           })
         | None => Constructor
         };
-      }
+      };
     | Undefined =>
       let. _ = otherwise(env, d);
       Indet;