Skip to main content

aranya_runtime/storage/
memory.rs

1use alloc::{boxed::Box, collections::BTreeMap, string::String, sync::Arc, vec::Vec};
2use core::ops::{Bound, Deref};
3
4use buggy::{Bug, BugExt as _, bug};
5use vec1::Vec1;
6
7use crate::{
8    Address, Checkpoint, CmdId, Command, Fact, FactIndex, FactPerspective, GraphId, Keys, Location,
9    MaxCut, Perspective, PolicyId, Prior, Priority, Query, QueryMut, Revertable, Segment,
10    SegmentIndex, Storage, StorageError, StorageProvider,
11};
12
13#[derive(Debug)]
14pub struct MemCommand {
15    priority: Priority,
16    id: CmdId,
17    parent: Prior<Address>,
18    policy: Option<Box<[u8]>>,
19    data: Box<[u8]>,
20    max_cut: MaxCut,
21}
22
23impl MemCommand {
24    fn from_cmd<C: Command>(command: &C, max_cut: MaxCut) -> Self {
25        let policy = command.policy().map(Box::from);
26
27        Self {
28            priority: command.priority(),
29            id: command.id(),
30            parent: command.parent(),
31            policy,
32            data: command.bytes().into(),
33            max_cut,
34        }
35    }
36}
37
38impl Command for MemCommand {
39    fn priority(&self) -> Priority {
40        self.priority.clone()
41    }
42
43    fn id(&self) -> CmdId {
44        self.id
45    }
46
47    fn parent(&self) -> Prior<Address> {
48        self.parent
49    }
50
51    fn policy(&self) -> Option<&[u8]> {
52        self.policy.as_deref()
53    }
54
55    fn bytes(&self) -> &[u8] {
56        &self.data
57    }
58
59    fn max_cut(&self) -> Result<MaxCut, Bug> {
60        Ok(self.max_cut)
61    }
62}
63
64#[derive(Default)]
65pub struct MemStorageProvider {
66    storage: BTreeMap<GraphId, MemStorage>,
67}
68
69impl MemStorageProvider {
70    pub const fn new() -> Self {
71        Self {
72            storage: BTreeMap::new(),
73        }
74    }
75}
76
77impl StorageProvider for MemStorageProvider {
78    type Perspective = MemPerspective;
79    type Storage = MemStorage;
80    type Segment = MemSegment;
81
82    fn new_perspective(&mut self, policy_id: PolicyId) -> Self::Perspective {
83        MemPerspective::new_unrooted(policy_id)
84    }
85
86    fn new_storage(
87        &mut self,
88        update: Self::Perspective,
89    ) -> Result<(GraphId, &mut Self::Storage), StorageError> {
90        use alloc::collections::btree_map::Entry;
91
92        let first = update
93            .commands
94            .first()
95            .ok_or(StorageError::EmptyPerspective)?;
96        let graph_id = GraphId::transmute(first.command.id);
97        let entry = match self.storage.entry(graph_id) {
98            Entry::Vacant(v) => v,
99            Entry::Occupied(_) => return Err(StorageError::StorageExists),
100        };
101
102        let mut storage = MemStorage::new();
103        let segment = storage.write(update)?;
104        storage.commit(segment)?;
105        Ok((graph_id, entry.insert(storage)))
106    }
107
108    fn get_storage(&mut self, graph: GraphId) -> Result<&mut Self::Storage, StorageError> {
109        self.storage
110            .get_mut(&graph)
111            .ok_or(StorageError::NoSuchStorage)
112    }
113
114    fn remove_storage(&mut self, graph: GraphId) -> Result<(), StorageError> {
115        self.storage
116            .remove(&graph)
117            .ok_or(StorageError::NoSuchStorage)?;
118
119        Ok(())
120    }
121
122    fn list_graph_ids(
123        &mut self,
124    ) -> Result<impl Iterator<Item = Result<GraphId, StorageError>>, StorageError> {
125        Ok(self.storage.keys().copied().map(Ok))
126    }
127}
128
129type FactMap = BTreeMap<Keys, Option<Box<[u8]>>>;
130type NamedFactMap = BTreeMap<String, FactMap>;
131
132pub struct MemStorage {
133    segments: Vec<MemSegment>,
134    commands: BTreeMap<CmdId, Location>,
135    head: Option<Location>,
136}
137
138impl MemStorage {
139    fn new() -> Self {
140        Self {
141            segments: Vec::new(),
142            commands: BTreeMap::new(),
143            head: None,
144        }
145    }
146
147    fn new_segment(
148        &mut self,
149        prior: Prior<Location>,
150        policy: PolicyId,
151        mut commands: Vec1<CommandData>,
152        facts: MemFactIndex,
153        max_cut: MaxCut,
154    ) -> Result<MemSegment, StorageError> {
155        let index = SegmentIndex(self.segments.len());
156
157        for (i, command) in commands.iter_mut().enumerate() {
158            command.command.max_cut = max_cut.checked_add(i).assume("must not overflow")?;
159        }
160
161        let segment = MemSegmentInner {
162            index,
163            prior,
164            policy,
165            commands,
166            facts,
167        };
168
169        let cell = MemSegment::from(segment);
170        self.segments.push(cell.clone());
171
172        Ok(cell)
173    }
174}
175
176impl Drop for MemStorage {
177    // Ensure the segments are dropped high to low, which helps avoid a stack
178    // overflow on dropping really long Arc chains.
179    fn drop(&mut self) {
180        while self.segments.pop().is_some() {}
181    }
182}
183
184impl Storage for MemStorage {
185    type Perspective = MemPerspective;
186    type Segment = MemSegment;
187    type FactIndex = MemFactIndex;
188    type FactPerspective = MemFactPerspective;
189
190    fn get_linear_perspective(&self, parent: Location) -> Result<Self::Perspective, StorageError> {
191        let segment = self.get_segment(parent)?;
192        let command = segment
193            .get_command(parent)
194            .ok_or(StorageError::CommandOutOfBounds(parent))?;
195        let parent_addr = command.address()?;
196
197        let policy = segment.policy;
198        let prior_facts: FactPerspectivePrior = if parent == segment.head_location()? {
199            segment.facts.clone().into()
200        } else {
201            let mut facts = MemFactPerspective::new(segment.facts.prior.clone().into());
202            for data in &segment.commands[..=segment.cmd_index(parent.max_cut)?] {
203                facts.apply_updates(&data.updates);
204            }
205            if facts.map.is_empty() {
206                facts.prior
207            } else {
208                facts.into()
209            }
210        };
211        let prior = Prior::Single(parent);
212        let parents = Prior::Single(parent_addr);
213
214        let max_cut = self
215            .get_segment(parent)?
216            .get_command(parent)
217            .assume("location must exist")?
218            .max_cut()?
219            .checked_add(1)
220            .assume("must not overflow")?;
221        let perspective = MemPerspective::new(prior, parents, policy, prior_facts, max_cut);
222
223        Ok(perspective)
224    }
225
226    fn get_fact_perspective(
227        &self,
228        location: Location,
229    ) -> Result<Self::FactPerspective, StorageError> {
230        let segment = self.get_segment(location)?;
231
232        if location == segment.head_location()?
233            || segment.commands.iter().all(|cmd| cmd.updates.is_empty())
234        {
235            return Ok(MemFactPerspective::new(segment.facts.clone().into()));
236        }
237
238        let mut facts = MemFactPerspective::new(segment.facts.prior.clone().into());
239        for data in &segment.commands[..=segment.cmd_index(location.max_cut)?] {
240            facts.apply_updates(&data.updates);
241        }
242
243        Ok(facts)
244    }
245
246    fn new_merge_perspective(
247        &self,
248        left: Location,
249        right: Location,
250        _last_common_ancestor: Location,
251        policy_id: PolicyId,
252        braid: MemFactIndex,
253    ) -> Result<Self::Perspective, StorageError> {
254        // TODO(jdygert): ensure braid belongs to this storage.
255        // TODO(jdygert): ensure braid ends at given command?
256
257        let left_segment = self.get_segment(left)?;
258        let left_policy_id = left_segment.policy;
259        let right_segment = self.get_segment(right)?;
260        let right_policy_id = right_segment.policy;
261
262        if (policy_id != left_policy_id) && (policy_id != right_policy_id) {
263            return Err(StorageError::PolicyMismatch);
264        }
265
266        let prior = Prior::Merge(left, right);
267
268        let left_command = left_segment
269            .get_command(left)
270            .ok_or(StorageError::CommandOutOfBounds(left))?;
271        let right_command = right_segment
272            .get_command(right)
273            .ok_or(StorageError::CommandOutOfBounds(right))?;
274        let parents = Prior::Merge(left_command.address()?, right_command.address()?);
275
276        let left_distance = left_command.max_cut()?;
277        let right_distance = right_command.max_cut()?;
278        let max_cut = left_distance
279            .max(right_distance)
280            .checked_add(1)
281            .assume("must not overflow")?;
282
283        let perspective = MemPerspective::new(prior, parents, policy_id, braid.into(), max_cut);
284
285        Ok(perspective)
286    }
287
288    fn get_segment(&self, location: Location) -> Result<MemSegment, StorageError> {
289        self.segments
290            .get(location.segment.0)
291            .ok_or(StorageError::SegmentOutOfBounds(location))
292            .cloned()
293    }
294
295    fn get_head(&self) -> Result<Location, StorageError> {
296        Ok(self.head.assume("storage has head after init")?)
297    }
298
299    fn write(&mut self, update: Self::Perspective) -> Result<Self::Segment, StorageError> {
300        let facts = self.write_facts(update.facts)?;
301
302        let commands: Vec1<CommandData> = update
303            .commands
304            .try_into()
305            .map_err(|_| StorageError::EmptyPerspective)?;
306
307        let segment_index = SegmentIndex(self.segments.len());
308
309        // Add the commands to the segment
310        for data in &commands {
311            let new_location = Location::new(segment_index, data.command.max_cut);
312            self.commands.insert(data.command.id(), new_location);
313        }
314
315        let segment =
316            self.new_segment(update.prior, update.policy, commands, facts, update.max_cut)?;
317
318        Ok(segment)
319    }
320
321    fn write_facts(
322        &mut self,
323        facts: Self::FactPerspective,
324    ) -> Result<Self::FactIndex, StorageError> {
325        let prior = match facts.prior {
326            FactPerspectivePrior::None => None,
327            FactPerspectivePrior::FactPerspective(prior) => Some(self.write_facts(*prior)?),
328            FactPerspectivePrior::FactIndex(prior) => Some(prior),
329        };
330        if facts.map.is_empty()
331            && let Some(prior) = prior
332        {
333            return Ok(prior);
334        }
335        Ok(MemFactIndex(Arc::new(MemFactsInner {
336            map: facts.map,
337            prior,
338        })))
339    }
340
341    fn commit(&mut self, segment: Self::Segment) -> Result<(), StorageError> {
342        // TODO(jdygert): ensure segment belongs to self?
343
344        if let Some(head) = self.head
345            && !self.is_ancestor(head, &segment)?
346        {
347            return Err(StorageError::HeadNotAncestor);
348        }
349
350        self.head = Some(segment.head_location()?);
351        Ok(())
352    }
353}
354
355#[derive(Clone, Debug)]
356pub struct MemFactIndex(Arc<MemFactsInner>);
357
358impl Deref for MemFactIndex {
359    type Target = MemFactsInner;
360    fn deref(&self) -> &Self::Target {
361        self.0.deref()
362    }
363}
364
365impl MemFactIndex {
366    #[cfg(all(test, feature = "graphviz"))]
367    fn name(&self) -> String {
368        format!("\"{:p}\"", Arc::as_ptr(&self.0))
369    }
370}
371
372#[derive(Debug)]
373pub struct MemFactsInner {
374    map: NamedFactMap,
375    prior: Option<MemFactIndex>,
376}
377
378pub(crate) fn find_prefixes<'m, 'p: 'm>(
379    map: &'m FactMap,
380    prefix: &'p [Box<[u8]>],
381) -> impl Iterator<Item = (&'m Keys, Option<&'m [u8]>)> + 'm {
382    map.range::<[Box<[u8]>], _>((Bound::Included(prefix), Bound::Unbounded))
383        .take_while(|(k, _)| k.starts_with(prefix))
384        .map(|(k, v)| (k, v.as_deref()))
385}
386
387impl FactIndex for MemFactIndex {}
388impl Query for MemFactIndex {
389    fn query(&self, name: &str, keys: &[Box<[u8]>]) -> Result<Option<Box<[u8]>>, StorageError> {
390        let mut prior = Some(self.deref());
391        while let Some(facts) = prior {
392            if let Some(slot) = facts.map.get(name).and_then(|m| m.get(keys)) {
393                return Ok(slot.clone());
394            }
395            prior = facts.prior.as_deref();
396        }
397        Ok(None)
398    }
399
400    type QueryIterator = Box<dyn Iterator<Item = Result<Fact, StorageError>>>;
401    fn query_prefix(
402        &self,
403        name: &str,
404        prefix: &[Box<[u8]>],
405    ) -> Result<Self::QueryIterator, StorageError> {
406        Ok(Box::from(
407            self.query_prefix_inner(name, prefix)
408                .into_iter()
409                // remove deleted facts
410                .filter_map(|(key, value)| Some(Ok(Fact { key, value: value? }))),
411        ))
412    }
413}
414
415impl MemFactIndex {
416    fn query_prefix_inner(&self, name: &str, prefix: &[Box<[u8]>]) -> FactMap {
417        let mut matches = BTreeMap::new();
418
419        let mut prior = Some(self.deref());
420        // walk backwards along fact indices
421        while let Some(facts) = prior {
422            if let Some(map) = facts.map.get(name) {
423                for (k, v) in find_prefixes(map, prefix) {
424                    // don't override, if we've already found the fact (including deletions)
425                    if !matches.contains_key(k) {
426                        matches.insert(k.clone(), v.map(Into::into));
427                    }
428                }
429            }
430            prior = facts.prior.as_deref();
431        }
432
433        matches
434    }
435}
436
437#[derive(Debug)]
438struct CommandData {
439    command: MemCommand,
440    updates: Vec<Update>,
441}
442
443#[derive(Debug)]
444pub struct MemSegmentInner {
445    index: SegmentIndex,
446    prior: Prior<Location>,
447    policy: PolicyId,
448    commands: Vec1<CommandData>,
449    facts: MemFactIndex,
450}
451
452#[derive(Clone, Debug)]
453pub struct MemSegment(Arc<MemSegmentInner>);
454
455impl Deref for MemSegment {
456    type Target = MemSegmentInner;
457
458    fn deref(&self) -> &Self::Target {
459        self.0.deref()
460    }
461}
462
463impl From<MemSegmentInner> for MemSegment {
464    fn from(segment: MemSegmentInner) -> Self {
465        Self(Arc::new(segment))
466    }
467}
468
469impl MemSegment {
470    fn cmd_index(&self, max_cut: MaxCut) -> Result<usize, StorageError> {
471        max_cut
472            .distance_from(self.shortest_max_cut())
473            .ok_or(StorageError::CommandOutOfBounds(Location::new(
474                self.index, max_cut,
475            )))
476    }
477}
478
479impl Segment for MemSegment {
480    type FactIndex = MemFactIndex;
481    type Command<'a> = &'a MemCommand;
482
483    fn index(&self) -> SegmentIndex {
484        self.index
485    }
486
487    fn head_id(&self) -> CmdId {
488        self.commands.last().command.id
489    }
490
491    fn first_location(&self) -> Location {
492        Location {
493            segment: self.index,
494            max_cut: self.commands.first().command.max_cut,
495        }
496    }
497
498    fn policy(&self) -> PolicyId {
499        self.policy
500    }
501
502    fn prior(&self) -> Prior<Location> {
503        self.prior
504    }
505
506    fn get_command(&self, location: Location) -> Option<&MemCommand> {
507        if location.segment != self.index {
508            return None;
509        }
510
511        self.commands
512            .get(self.cmd_index(location.max_cut).ok()?)
513            .map(|d| &d.command)
514    }
515
516    fn longest_max_cut(&self) -> Result<MaxCut, StorageError> {
517        Ok(self.commands.last().command.max_cut)
518    }
519
520    fn shortest_max_cut(&self) -> MaxCut {
521        self.commands.first().command.max_cut
522    }
523
524    fn skip_list(&self) -> &[Location] {
525        &[]
526    }
527
528    fn facts(&self) -> Result<Self::FactIndex, StorageError> {
529        Ok(self.facts.clone())
530    }
531}
532
533type Update = (String, Keys, Option<Box<[u8]>>);
534
535#[derive(Debug)]
536pub struct MemPerspective {
537    prior: Prior<Location>,
538    parents: Prior<Address>,
539    policy: PolicyId,
540    facts: MemFactPerspective,
541    commands: Vec<CommandData>,
542    current_updates: Vec<Update>,
543    max_cut: MaxCut,
544}
545
546#[derive(Debug)]
547enum FactPerspectivePrior {
548    None,
549    FactPerspective(Box<MemFactPerspective>),
550    FactIndex(MemFactIndex),
551}
552
553impl From<MemFactIndex> for FactPerspectivePrior {
554    fn from(value: MemFactIndex) -> Self {
555        Self::FactIndex(value)
556    }
557}
558
559impl From<Option<MemFactIndex>> for FactPerspectivePrior {
560    fn from(value: Option<MemFactIndex>) -> Self {
561        value.map_or(Self::None, Self::FactIndex)
562    }
563}
564
565impl From<MemFactPerspective> for FactPerspectivePrior {
566    fn from(value: MemFactPerspective) -> Self {
567        Self::FactPerspective(Box::new(value))
568    }
569}
570
571#[derive(Debug)]
572pub struct MemFactPerspective {
573    map: NamedFactMap,
574    prior: FactPerspectivePrior,
575}
576
577impl MemFactPerspective {
578    fn new(prior_facts: FactPerspectivePrior) -> Self {
579        Self {
580            map: NamedFactMap::new(),
581            prior: prior_facts,
582        }
583    }
584
585    fn clear(&mut self) {
586        self.map.clear();
587    }
588
589    fn apply_updates(&mut self, updates: &[Update]) {
590        for (name, key, value) in updates {
591            self.map
592                .entry(name.clone())
593                .or_default()
594                .insert(key.clone(), value.clone());
595        }
596    }
597}
598
599impl MemPerspective {
600    fn new(
601        prior: Prior<Location>,
602        parents: Prior<Address>,
603        policy: PolicyId,
604        prior_facts: FactPerspectivePrior,
605        max_cut: MaxCut,
606    ) -> Self {
607        Self {
608            prior,
609            parents,
610            policy,
611            facts: MemFactPerspective::new(prior_facts),
612            commands: Vec::new(),
613            current_updates: Vec::new(),
614            max_cut,
615        }
616    }
617
618    fn new_unrooted(policy: PolicyId) -> Self {
619        Self {
620            prior: Prior::None,
621            parents: Prior::None,
622            policy,
623            facts: MemFactPerspective::new(FactPerspectivePrior::None),
624            commands: Vec::new(),
625            current_updates: Vec::new(),
626            max_cut: MaxCut(0),
627        }
628    }
629}
630
631impl Revertable for MemPerspective {
632    fn checkpoint(&self) -> Checkpoint {
633        Checkpoint {
634            index: self.commands.len(),
635        }
636    }
637
638    fn revert(&mut self, checkpoint: Checkpoint) -> Result<(), Bug> {
639        if checkpoint.index == self.commands.len() {
640            return Ok(());
641        }
642
643        if checkpoint.index > self.commands.len() {
644            bug!(
645                "A checkpoint's index should always be less than or equal to the length of a perspective's command history!"
646            );
647        }
648
649        self.commands.truncate(checkpoint.index);
650        self.facts.clear();
651        self.current_updates.clear();
652        for data in &self.commands {
653            self.facts.apply_updates(&data.updates);
654        }
655
656        Ok(())
657    }
658}
659
660impl Perspective for MemPerspective {
661    fn add_command(&mut self, command: &impl Command) -> Result<usize, StorageError> {
662        if command.parent() != self.head_address()? {
663            return Err(StorageError::PerspectiveHeadMismatch);
664        }
665
666        let entry = CommandData {
667            command: MemCommand::from_cmd(command, self.head_address()?.next_max_cut()?),
668            updates: core::mem::take(&mut self.current_updates),
669        };
670        self.commands.push(entry);
671        Ok(self.commands.len())
672    }
673
674    fn policy(&self) -> PolicyId {
675        self.policy
676    }
677
678    fn includes(&self, id: CmdId) -> bool {
679        self.commands.iter().any(|cmd| cmd.command.id == id)
680    }
681
682    fn head_address(&self) -> Result<Prior<Address>, Bug> {
683        Ok(if let Some(last) = self.commands.last() {
684            Prior::Single(last.command.address()?)
685        } else {
686            self.parents
687        })
688    }
689}
690
691impl FactPerspective for MemPerspective {}
692
693impl Query for MemPerspective {
694    fn query(&self, name: &str, keys: &[Box<[u8]>]) -> Result<Option<Box<[u8]>>, StorageError> {
695        self.facts.query(name, keys)
696    }
697
698    type QueryIterator = <MemFactPerspective as Query>::QueryIterator;
699    fn query_prefix(
700        &self,
701        name: &str,
702        prefix: &[Box<[u8]>],
703    ) -> Result<Self::QueryIterator, StorageError> {
704        self.facts.query_prefix(name, prefix)
705    }
706}
707
708impl QueryMut for MemPerspective {
709    fn insert(&mut self, name: String, keys: Keys, value: Box<[u8]>) {
710        self.facts.insert(name.clone(), keys.clone(), value.clone());
711        self.current_updates.push((name, keys, Some(value)));
712    }
713
714    fn delete(&mut self, name: String, keys: Keys) {
715        self.facts.delete(name.clone(), keys.clone());
716        self.current_updates.push((name, keys, None));
717    }
718}
719
720impl MemFactPerspective {
721    fn query_prefix_inner(&self, name: &str, prefix: &[Box<[u8]>]) -> FactMap {
722        let map = self.map.get(name);
723        let mut matches = match &self.prior {
724            FactPerspectivePrior::None => BTreeMap::new(),
725            FactPerspectivePrior::FactPerspective(fp) => fp.query_prefix_inner(name, prefix),
726            FactPerspectivePrior::FactIndex(fi) => fi.query_prefix_inner(name, prefix),
727        };
728        if let Some(map) = map {
729            for (k, v) in find_prefixes(map, prefix) {
730                // overwrite "earlier" facts
731                matches.insert(k.clone(), v.map(Into::into));
732            }
733        }
734        matches
735    }
736}
737
738impl FactPerspective for MemFactPerspective {}
739
740impl Query for MemFactPerspective {
741    fn query(&self, name: &str, keys: &[Box<[u8]>]) -> Result<Option<Box<[u8]>>, StorageError> {
742        if let Some(wrapped) = self.map.get(name).and_then(|m| m.get(keys)) {
743            return Ok(wrapped.as_deref().map(Box::from));
744        }
745        match &self.prior {
746            FactPerspectivePrior::None => Ok(None),
747            FactPerspectivePrior::FactPerspective(prior) => prior.query(name, keys),
748            FactPerspectivePrior::FactIndex(prior) => prior.query(name, keys),
749        }
750    }
751
752    type QueryIterator = Box<dyn Iterator<Item = Result<Fact, StorageError>>>;
753    fn query_prefix(
754        &self,
755        name: &str,
756        prefix: &[Box<[u8]>],
757    ) -> Result<Self::QueryIterator, StorageError> {
758        Ok(Box::from(
759            self.query_prefix_inner(name, prefix)
760                .into_iter()
761                // remove deleted facts
762                .filter_map(|(key, value)| Some(Ok(Fact { key, value: value? }))),
763        ))
764    }
765}
766
767impl QueryMut for MemFactPerspective {
768    fn insert(&mut self, name: String, keys: Keys, value: Box<[u8]>) {
769        self.map.entry(name).or_default().insert(keys, Some(value));
770    }
771
772    fn delete(&mut self, name: String, keys: Keys) {
773        self.map.entry(name).or_default().insert(keys, None);
774    }
775}
776
777#[cfg(all(test, feature = "graphviz"))]
778pub mod graphviz {
779    #![allow(clippy::unwrap_used)]
780
781    use std::{fs::File, io::BufWriter};
782
783    use dot_writer::{Attributes as _, DotWriter, Style};
784
785    #[allow(clippy::wildcard_imports)]
786    use super::*;
787    use crate::testing::short_b58;
788
789    fn loc(location: impl Into<Location>) -> String {
790        let location = location.into();
791        format!("\"{}:{}\"", location.segment, location.max_cut)
792    }
793
794    fn get_seq(p: &MemFactIndex) -> &str {
795        if let Some(Some(seq)) = p.map.get("seq").and_then(|m| m.get(&Keys::default())) {
796            std::str::from_utf8(seq).unwrap()
797        } else {
798            ""
799        }
800    }
801
802    fn dotwrite(storage: &MemStorage, out: &mut DotWriter<'_>) {
803        let mut graph = out.digraph();
804        graph
805            .graph_attributes()
806            .set("compound", "true", false)
807            .set("rankdir", "RL", false)
808            .set_style(Style::Filled)
809            .set("color", "grey", false);
810        graph
811            .node_attributes()
812            .set("shape", "square", false)
813            .set_style(Style::Filled)
814            .set("color", "lightgrey", false);
815
816        let mut seen_facts = std::collections::HashMap::new();
817        let mut external_facts = Vec::new();
818
819        for segment in &storage.segments {
820            let mut cluster = graph.cluster();
821            match segment.prior {
822                Prior::None => {
823                    cluster.graph_attributes().set("color", "green", false);
824                }
825                Prior::Single(..) => {}
826                Prior::Merge(..) => {
827                    cluster.graph_attributes().set("color", "crimson", false);
828                }
829            }
830
831            // Draw commands and edges between commands within the segment.
832            for (i, cmd) in segment.commands.iter().enumerate() {
833                {
834                    let mut node = cluster.node_named(loc((segment.index, cmd.command.max_cut)));
835                    node.set_label(&short_b58(cmd.command.id));
836                    match cmd.command.parent {
837                        Prior::None => {
838                            node.set("shape", "house", false);
839                        }
840                        Prior::Single(..) => {}
841                        Prior::Merge(..) => {
842                            node.set("shape", "hexagon", false);
843                        }
844                    }
845                }
846                if i > 0 {
847                    let previous = cmd.command.max_cut.decremented().expect("i must be > 0");
848                    cluster.edge(
849                        loc((segment.index, cmd.command.max_cut)),
850                        loc((segment.index, previous)),
851                    );
852                }
853            }
854
855            // Draw edges to previous segments.
856            let first = loc(segment.first_location());
857            for p in segment.prior() {
858                cluster.edge(&first, loc(p));
859            }
860
861            // Draw fact index for this segment.
862            let curr = segment.facts.name();
863            cluster
864                .node_named(curr.clone())
865                .set_label(get_seq(&segment.facts))
866                .set("shape", "cylinder", false)
867                .set("color", "black", false)
868                .set("style", "solid", false);
869            cluster
870                .edge(loc(segment.head_location().unwrap()), &curr)
871                .attributes()
872                .set("color", "red", false);
873
874            seen_facts.insert(curr, segment.facts.clone());
875            // Make sure prior facts of fact index will get processed later.
876            let mut node = &segment.facts;
877            while let Some(prior) = &node.prior {
878                node = prior;
879                let name = node.name();
880                if seen_facts.insert(name, node.clone()).is_some() {
881                    break;
882                }
883                external_facts.push(node.clone());
884            }
885        }
886
887        graph
888            .node_attributes()
889            .set("shape", "cylinder", false)
890            .set("color", "black", false)
891            .set("style", "solid", false);
892
893        for fact in external_facts {
894            // Draw nodes for fact indices not directly associated with a segment.
895            graph.node_named(fact.name()).set_label(get_seq(&fact));
896
897            // Draw edge to prior facts.
898            if let Some(prior) = &fact.prior {
899                graph
900                    .edge(fact.name(), prior.name())
901                    .attributes()
902                    .set("color", "blue", false);
903            }
904        }
905
906        // Draw edges to prior facts for fact indices in segments.
907        for segment in &storage.segments {
908            if let Some(prior) = &segment.facts.prior {
909                graph
910                    .edge(segment.facts.name(), prior.name())
911                    .attributes()
912                    .set("color", "blue", false);
913            }
914        }
915
916        // Draw HEAD indicator.
917        graph.node_named("HEAD").set("shape", "none", false);
918        graph.edge("HEAD", loc(storage.get_head().unwrap()));
919    }
920
921    pub fn dot(storage: &MemStorage, name: &str) {
922        std::fs::create_dir_all(".ignore").unwrap();
923        dotwrite(
924            storage,
925            &mut DotWriter::from(&mut BufWriter::new(
926                File::create(format!(".ignore/{name}.dot")).unwrap(),
927            )),
928        );
929    }
930}
931
932#[cfg(test)]
933mod test {
934    use super::*;
935    use crate::testing::dsl::{StorageBackend, test_suite};
936
937    #[test]
938    fn test_query_prefix() {
939        let mut graph = MemStorage::new();
940        let mut fp = MemFactPerspective::new(FactPerspectivePrior::None);
941
942        let name = "x";
943
944        let keys: &[&[&str]] = &[
945            &["aa", "xy", "123"],
946            &["aa", "xz", "123"],
947            &["bb", "ccc"],
948            &["bc", ""],
949        ];
950        let keys: Vec<Keys> = keys
951            .iter()
952            .map(|ks| ks.iter().map(|k| k.as_bytes()).collect())
953            .collect();
954
955        for ks in &keys {
956            fp.insert(
957                name.into(),
958                ks.clone(),
959                format!("{ks:?}").into_bytes().into(),
960            );
961        }
962        let facts = graph.write_facts(fp).unwrap();
963
964        let prefixes: &[&[&str]] = &[
965            &["aa", "xy", "12"],
966            &["aa", "xy"],
967            &["aa", "xz"],
968            &["aa", "x"],
969            &["bb", ""],
970            &["bb", "ccc"],
971            &["bc", ""],
972            &["bc", "", ""],
973        ];
974
975        for prefix in prefixes {
976            let prefix: Keys = prefix.iter().map(|k| k.as_bytes()).collect();
977            let found: Vec<_> = facts.query_prefix(name, &prefix).unwrap().collect();
978            let mut expected: Vec<_> = keys.iter().filter(|k| k.starts_with(&prefix)).collect();
979            expected.sort();
980            assert_eq!(found.len(), expected.len());
981            for (a, b) in std::iter::zip(found, expected) {
982                let a = a.unwrap();
983                assert_eq!(&a.key, b);
984                assert_eq!(a.value.as_ref(), format!("{b:?}").as_bytes());
985            }
986        }
987    }
988
989    struct MemBackend;
990    impl StorageBackend for MemBackend {
991        type StorageProvider = MemStorageProvider;
992
993        fn provider(&mut self, _client_id: u64) -> Self::StorageProvider {
994            MemStorageProvider::new()
995        }
996    }
997    test_suite!(|| MemBackend);
998}