mangle_factstore/
tablestore.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 std::cell::RefCell;
16use std::collections::HashMap;
17use std::path::Path;
18
19use crate::ast;
20use crate::Bump;
21use crate::FactStore;
22use crate::ReadOnlyFactStore;
23use crate::{anyhow, Result};
24
25#[derive(Clone)]
26pub enum TableConfig<'a> {
27    InMemory,
28    RowFile(&'a Path),
29}
30
31pub type TableStoreSchema<'a> = HashMap<&'a ast::PredicateSym<'a>, TableConfig<'a>>;
32
33pub struct TableStoreImpl<'a> {
34    schema: &'a TableStoreSchema<'a>,
35    bump: Bump,
36    tables: RefCell<HashMap<&'a ast::PredicateSym<'a>, Vec<&'a ast::Atom<'a>>>>,
37}
38
39impl<'a> ReadOnlyFactStore<'a> for TableStoreImpl<'a> {
40    fn contains(&'a self, fact: &ast::Atom) -> Result<bool> {
41        if let Some(table) = self.tables.borrow().get(&fact.sym) {
42            return Ok(table.contains(&fact));
43        }
44        Err(anyhow!("unknown predicate {}", fact.sym.name))
45    }
46
47    fn get(
48        &'a self,
49        query: &ast::Atom,
50        mut cb: impl FnMut(&'a ast::Atom<'a>) -> anyhow::Result<()>,
51    ) -> anyhow::Result<()> {
52        if let Some(table) = self.tables.borrow().get(&query.sym) {
53            for fact in table {
54                if fact.matches(query.args) {
55                    cb(fact)?;
56                }
57            }
58            return Ok(());
59        }
60        Err(anyhow!("unknown predicate {}", query.sym.name))
61    }
62
63    fn list_predicates(&'a self, mut cb: impl FnMut(&'a ast::PredicateSym)) {
64        for pred in self.tables.borrow().keys() {
65            cb(pred);
66        }
67    }
68
69    fn estimate_fact_count(&self) -> u32 {
70        self.tables
71            .borrow()
72            .values()
73            .fold(0, |x, y| x + y.len() as u32)
74    }
75}
76
77impl<'a> FactStore<'a> for TableStoreImpl<'a> {
78    fn add(&'a self, fact: &ast::Atom) -> Result<bool> {
79        {
80            if let Some(table) = self.tables.borrow().get(&fact.sym) {
81                if table.contains(&fact) {
82                    return Ok(false);
83                }
84            } else {
85                return Err(anyhow!("no table for {:?}", fact.sym.name));
86            }
87        }
88        // Consider checking that `fact` is, in fact, a fact.
89        let fact = ast::copy_atom(&self.bump, fact);
90        self.tables
91            .borrow_mut()
92            .get_mut(&fact.sym)
93            .unwrap()
94            .push(fact);
95        Ok(true)
96    }
97
98    fn merge<'other, S>(&'a self, _store: &'other S)
99    where
100        S: crate::ReadOnlyFactStore<'other>,
101    {
102        todo!()
103    }
104}
105
106impl<'a> TableStoreImpl<'a> {
107    pub fn new(schema: &'a TableStoreSchema<'a>) -> Self {
108        let mut tables = HashMap::new();
109        for entry in schema.keys() {
110            tables.insert(*entry, vec![]);
111        }
112        TableStoreImpl {
113            schema,
114            bump: Bump::new(),
115            tables: RefCell::new(tables),
116        }
117    }
118}