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

mangle_common/
tablestore.rs
