mangle_factstore/
lib.rs

1// Copyright 2024 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use anyhow::{anyhow, Result};
16use bumpalo::Bump;
17use mangle_ast as ast;
18
19mod tablestore;
20pub use tablestore::{TableConfig, TableStoreImpl, TableStoreSchema};
21
22/// Lifetime 'a is used for data held by this store.
23pub trait ReadOnlyFactStore<'a> {
24    fn contains(&'a self, fact: &ast::Atom) -> Result<bool>;
25
26    // Invokes cb for fact that matches query.
27    fn get(
28        &'a self,
29        query: &ast::Atom,
30        cb: impl FnMut(&'a ast::Atom<'a>) -> Result<()>,
31    ) -> Result<()>;
32
33    //fn get<F>(&'a self, query: &ast::Atom, cb: F) -> Result<()>
34    //where
35    //    F: FnMut(&'a ast::Atom<'a>) -> Result<()>;
36
37    // Invokes cb for every predicate available in this store.
38    // It would be nice to use `impl Iterator` here.
39    fn list_predicates(&'a self, cb: impl FnMut(&'a ast::PredicateSym));
40
41    // Returns approximae number of facts.
42    fn estimate_fact_count(&self) -> u32;
43}
44
45pub trait FactStore<'a>: ReadOnlyFactStore<'a> {
46    /// Returns true if fact did not exist before.
47    /// The fact is copied.
48    fn add(&'a self, fact: &ast::Atom) -> Result<bool>;
49
50    /// Adds all facts from given store.
51    fn merge<'other, S>(&'a self, store: &'other S)
52    where
53        S: ReadOnlyFactStore<'other>;
54}
55
56/// Constructs a query (allocated in bump).
57/// `pred.arity` must be present.
58/// TODO: move to ast
59/// TODO: make an allocator interface that has slice with all _ arguments.
60fn new_query<'b>(bump: &'b Bump, pred: &'b ast::PredicateSym) -> &'b ast::Atom<'b> {
61    let var = &*bump.alloc(ast::BaseTerm::Variable("_"));
62    let mut args = vec![];
63    for _ in 0..pred.arity.unwrap() {
64        args.push(var);
65    }
66    let args = &*bump.alloc_slice_copy(&args);
67
68    bump.alloc(ast::Atom { sym: *pred, args })
69}
70
71pub fn get_all_facts<'a, S>(
72    store: &'a S,
73    mut cb: impl FnMut(&'a ast::Atom<'a>) -> Result<()>,
74) -> Result<()>
75where
76    S: ReadOnlyFactStore<'a> + 'a,
77{
78    let bump = Bump::new();
79    let mut preds = vec![];
80    store.list_predicates(|pred: &ast::PredicateSym| {
81        preds.push(pred);
82    });
83    for pred in preds {
84        store.get(new_query(&bump, pred), &mut cb)?;
85    }
86    Ok(())
87}
88
89#[cfg(test)]
90mod test {
91    use std::{cell::RefCell, collections::HashSet};
92
93    use super::*;
94
95    static TEST_ATOM: ast::Atom = ast::Atom {
96        sym: ast::PredicateSym {
97            name: "foo",
98            arity: Some(1),
99        },
100        args: &[&ast::BaseTerm::Const(ast::Const::String("bar"))],
101    };
102
103    struct TestStore<'a> {
104        bump: &'a Bump,
105        facts: RefCell<Vec<&'a ast::Atom<'a>>>,
106    }
107
108    impl<'a> ReadOnlyFactStore<'a> for TestStore<'a> {
109        fn contains<'store>(&'store self, fact: &ast::Atom) -> Result<bool> {
110            Ok(self.facts.borrow().iter().any(|x| *x == fact))
111        }
112
113        fn get(
114            &'a self,
115            query: &ast::Atom,
116            mut cb: impl FnMut(&'a ast::Atom<'a>) -> Result<()>,
117        ) -> Result<()> {
118            for fact in self.facts.borrow().iter() {
119                // TODO matches
120                if fact.sym == query.sym {
121                    cb(fact)?;
122                }
123            }
124            Ok(())
125        }
126
127        fn list_predicates(&'a self, mut cb: impl FnMut(&'a ast::PredicateSym)) {
128            let mut seen = HashSet::new();
129            for fact in self.facts.borrow().iter() {
130                let pred = &fact.sym;
131                if !seen.contains(pred) {
132                    seen.insert(pred);
133                    cb(pred)
134                }
135            }
136        }
137
138        fn estimate_fact_count(&self) -> u32 {
139            self.facts.borrow().len().try_into().unwrap()
140        }
141    }
142
143    impl<'a> FactStore<'a> for TestStore<'a> {
144        fn add(&'a self, fact: &ast::Atom) -> Result<bool> {
145            if self.contains(fact)? {
146                return Ok(false);
147            }
148            self.facts
149                .borrow_mut()
150                .push(ast::copy_atom(&self.bump, fact));
151            Ok(true)
152        }
153
154        fn merge<'other, S>(&'a self, store: &'other S)
155        where
156            S: ReadOnlyFactStore<'other>,
157        {
158            let _ = get_all_facts(store, move |fact| {
159                let atom = ast::copy_atom(&self.bump, fact);
160                self.facts.borrow_mut().push(atom);
161                Ok(())
162            });
163        }
164    }
165
166    #[test]
167    fn test_get_factsa() {
168        let bump = Bump::new();
169        let simple = TestStore {
170            bump: &bump,
171            facts: RefCell::new(vec![]),
172        };
173        assert!(!simple.contains(&TEST_ATOM).unwrap());
174        assert!(simple.add(&TEST_ATOM).unwrap());
175    }
176}