probl_engine/
world.rs

1//! Worlds, and merging the ones that have become identical.
2
3use crate::analytic::Constraints;
4use crate::error::RuntimeError;
5use crate::value::Value;
6use crate::weight::Weight;
7use probl_sema::SlotSet;
8use probl_sema::ir::SlotId;
9use rustc_hash::{FxHashMap, FxHasher};
10use std::collections::BTreeSet;
11use std::collections::hash_map::Entry;
12use std::hash::{Hash, Hasher};
13use std::sync::Arc;
14
15/// A returned value, the weight of the world that returned it, and the
16/// restrictions on analytic draws that its caller can see.
17pub type Returned = (Value, Weight, Constraints);
18
19/// One possible state of the program: the variables of the current frame,
20/// and how likely it is (including every observation so far).
21#[derive(Clone, Debug)]
22pub struct World {
23    pub slots: Vec<Value>,
24    pub constraints: Constraints,
25    /// Latents passed into this call: restrictions on these remain observable
26    /// by its caller even after every local alias dies. Constant per frame.
27    pub inherited: Arc<BTreeSet<u64>>,
28    pub weight: Weight,
29    /// When sampling, the run this world is (docs/semantics.md, section 14).
30    pub run: u32,
31    /// When enumerating, how many densities the weight includes: weights
32    /// with different numbers are in different units, so they don't add up.
33    pub densities: u32,
34}
35
36impl World {
37    pub fn scaled(mut self, factor: f64) -> World {
38        self.weight = self.weight.scale(factor);
39        self
40    }
41}
42
43pub fn total_weight(worlds: &[World]) -> Weight {
44    Weight::sum(worlds.iter().map(|w| w.weight))
45}
46
47/// Worlds leaving a statement, grouped by how they left it.
48#[derive(Debug, Default)]
49pub struct Flow {
50    /// Continue with the next statement.
51    pub next: Vec<World>,
52    pub broke: Vec<World>,
53    pub continued: Vec<World>,
54    /// Left the function: the returned value and the world's weight.
55    pub returned: Vec<Returned>,
56    /// Faulted inside a `try` that may catch the fault: each world as it
57    /// was when its statement began, with its fault, on its way to the
58    /// `catch`.
59    pub faulted: Vec<(World, RuntimeError)>,
60}
61
62impl Flow {
63    pub fn next(worlds: Vec<World>) -> Flow {
64        Flow {
65            next: worlds,
66            ..Flow::default()
67        }
68    }
69
70    pub fn join(&mut self, other: Flow) {
71        self.next.extend(other.next);
72        self.broke.extend(other.broke);
73        self.continued.extend(other.continued);
74        self.returned.extend(other.returned);
75        self.faulted.extend(other.faulted);
76    }
77}
78
79/// Clear these slots: nobody will read them again.
80pub fn clear(worlds: &mut [World], slots: &[SlotId]) {
81    if slots.is_empty() {
82        return;
83    }
84    for w in worlds {
85        for &s in slots {
86            w.slots[s as usize] = Value::Dead;
87        }
88    }
89}
90
91/// Clear every slot that isn't live, in worlds that jumped (with `break` or
92/// `continue`) past statements that would have cleared some of them.
93pub fn clear_dead(worlds: &mut [World], live: &SlotSet) {
94    let Some(n) = worlds.first().map(|w| w.slots.len()) else {
95        return;
96    };
97    let dead: Vec<SlotId> = live.iter_missing(n).collect();
98    clear(worlds, &dead);
99}
100
101/// The live slots among a frame's `n`.
102pub fn live_slots(live: &SlotSet, n: usize) -> Vec<usize> {
103    live.iter().map(|s| s as usize).filter(|&s| s < n).collect()
104}
105
106/// A hash of a world's live slots, as merging computes it.
107pub fn state_hash(w: &World, live: &[usize]) -> u64 {
108    let mut hasher = FxHasher::default();
109    for &i in live {
110        w.slots[i].hash(&mut hasher);
111    }
112    w.constraints.hash(&mut hasher);
113    hasher.finish()
114}
115
116/// If `enabled`, merge the worlds that agree on the live slots, adding up
117/// their weights, and keeping the order of first appearance. Dead slots
118/// don't matter, cleared or not.
119pub fn merge(mut worlds: Vec<World>, live: &SlotSet, enabled: bool) -> Vec<World> {
120    for w in &mut worlds {
121        // Most worlds have no restrictions: don't even iterate.
122        if w.constraints.is_empty() || w.constraints.keys().all(|id| w.inherited.contains(id)) {
123            continue;
124        }
125        let mut ids = (*w.inherited).clone();
126        for slot in live.iter() {
127            crate::analytic::collect_ids(&w.slots[slot as usize], &mut ids);
128        }
129        crate::analytic::collect_form_ids(&w.constraints, &mut ids);
130        Arc::make_mut(&mut w.constraints).retain(|id, _| ids.contains(id));
131    }
132    if !enabled || worlds.len() < 2 {
133        return worlds;
134    }
135    let n = worlds[0].slots.len();
136    let live = live_slots(live, n);
137    let mut out: Vec<World> = Vec::with_capacity(worlds.len());
138    // The first world kept with each hash, and after each world, the next
139    // one with the same hash (only when different worlds' hashes collide).
140    let mut first: FxHashMap<u64, usize> = FxHashMap::with_capacity_and_hasher(worlds.len(), Default::default());
141    let mut next: Vec<usize> = Vec::with_capacity(worlds.len());
142    const NONE: usize = usize::MAX;
143    for w in worlds {
144        let same = |kept: &World| {
145            kept.densities == w.densities
146                && kept.constraints == w.constraints
147                && live.iter().all(|&i| kept.slots[i] == w.slots[i])
148        };
149        let found = match first.entry(state_hash(&w, &live)) {
150            Entry::Vacant(entry) => {
151                entry.insert(out.len());
152                None
153            }
154            Entry::Occupied(entry) => {
155                let mut i = *entry.get();
156                loop {
157                    if same(&out[i]) {
158                        break Some(i);
159                    }
160                    if next[i] == NONE {
161                        next[i] = out.len();
162                        break None;
163                    }
164                    i = next[i];
165                }
166            }
167        };
168        match found {
169            Some(i) => out[i].weight += w.weight,
170            None => {
171                next.push(NONE);
172                out.push(w);
173            }
174        }
175    }
176    out
177}
178
179/// Merge returned values only when their caller-visible restrictions also agree.
180pub fn merge_values(pairs: Vec<Returned>) -> Vec<Returned> {
181    let mut out: Vec<Returned> = Vec::with_capacity(pairs.len());
182    let mut index: FxHashMap<(Value, Constraints), usize> = FxHashMap::default();
183    for (v, w, constraints) in pairs {
184        match index.entry((v.clone(), constraints.clone())) {
185            Entry::Occupied(entry) => out[*entry.get()].1 += w,
186            Entry::Vacant(entry) => {
187                entry.insert(out.len());
188                out.push((v, w, constraints));
189            }
190        }
191    }
192    out
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198
199    fn world(slots: Vec<Value>) -> World {
200        World {
201            slots,
202            constraints: Default::default(),
203            inherited: Default::default(),
204            weight: Weight::ONE,
205            run: 0,
206            densities: 0,
207        }
208    }
209
210    #[test]
211    fn worlds_merge_on_their_live_slots() {
212        let mut live = SlotSet::with_capacity(3);
213        live.insert(0);
214        // Slot 1 is dead: worlds that differ only there are the same world.
215        let worlds = vec![
216            world(vec![Value::Int(1.into()), Value::Int(7.into()), Value::Dead]),
217            world(vec![Value::Int(2.into()), Value::Int(7.into()), Value::Dead]),
218            world(vec![
219                Value::Int(1.into()),
220                Value::list(vec![Value::Int(8.into())]),
221                Value::Dead,
222            ]),
223        ];
224        let mut merged = merge(worlds, &live, true);
225        assert_eq!(merged.len(), 2);
226        assert_eq!(merged[0].slots[0], Value::Int(1.into()));
227        assert_eq!(merged[0].weight.to_f64(), 2.0);
228        assert_eq!(merged[1].weight.to_f64(), 1.0);
229        clear_dead(&mut merged, &live);
230        assert!(merged.iter().all(|w| w.slots[1] == Value::Dead));
231        assert_eq!(merged[1].slots[0], Value::Int(2.into()));
232    }
233}
probl_engine/world.rs

probl_engine/
world.rs
