mangle_analysis/
simple_program.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 crate::{ast, PredicateSet, Program, StratifiedProgram};
16use bumpalo::Bump;
17use std::collections::HashMap;
18
19/// An implementation of the `Program` trait.
20pub struct SimpleProgram<'p> {
21    pub bump: &'p Bump,
22    pub ext_preds: Vec<ast::PredicateSym<'p>>,
23    pub rules: HashMap<ast::PredicateSym<'p>, Vec<&'p ast::Clause<'p>>>,
24}
25
26/// An implementation of the `StratifiedProgram` trait.
27/// This can be obtained through SimpleProgram::stratify.
28pub struct SimpleStratifiedProgram<'p> {
29    program: SimpleProgram<'p>,
30    strata: &'p [&'p PredicateSet<'p>],
31}
32
33impl<'p> Program<'p> for SimpleProgram<'p> {
34    fn extensional_preds(&'p self) -> impl Iterator<Item = &'p ast::PredicateSym<'p>> {
35        self.ext_preds.iter()
36    }
37
38    fn intensional_preds(&'p self) -> impl Iterator<Item = &'p ast::PredicateSym<'p>> {
39        self.rules.keys()
40    }
41
42    fn rules(
43        &'p self,
44        sym: &'p ast::PredicateSym<'p>,
45    ) -> impl Iterator<Item = &'p ast::Clause<'p>> {
46        self.rules.get(sym).unwrap().iter().copied()
47    }
48}
49
50impl<'p> SimpleProgram<'p> {
51    pub fn add_clause(&mut self, clause: &ast::Clause) {
52        let clause = ast::copy_clause(self.bump, clause);
53        let sym = clause.head.sym;
54        use std::collections::hash_map::Entry;
55        match self.rules.entry(sym) {
56            Entry::Occupied(mut v) => v.get_mut().push(clause),
57            Entry::Vacant(v) => {
58                v.insert(vec![clause]);
59            }
60        }
61    }
62
63    /// Produces a StratifiedProgram with given set of layers.
64    /// TODO: write analysis that computes the set of layers.
65    pub fn stratify(
66        self,
67        strata: impl Iterator<Item = &'p PredicateSet<'p>>,
68    ) -> SimpleStratifiedProgram<'p> {
69        let mut layers = vec![];
70        for layer in strata {
71            layers.push(*self.bump.alloc(layer))
72        }
73        let strata = &*self.bump.alloc_slice_copy(&layers);
74        SimpleStratifiedProgram {
75            program: self,
76            strata,
77        }
78    }
79}
80
81impl<'p> Program<'p> for SimpleStratifiedProgram<'p> {
82    fn extensional_preds(&'p self) -> impl Iterator<Item = &'p ast::PredicateSym<'p>> {
83        self.program.extensional_preds()
84    }
85
86    fn intensional_preds(&'p self) -> impl Iterator<Item = &'p ast::PredicateSym<'p>> {
87        self.program.intensional_preds()
88    }
89
90    fn rules(
91        &'p self,
92        sym: &'p ast::PredicateSym<'p>,
93    ) -> impl Iterator<Item = &'p ast::Clause<'p>> {
94        self.program.rules(sym)
95    }
96}
97
98impl<'p> StratifiedProgram<'p> for SimpleStratifiedProgram<'p> {
99    fn strata(&'p self) -> impl Iterator<Item = &'p PredicateSet<'p>> {
100        self.strata.iter().copied()
101    }
102
103    fn pred_to_index(&'p self, sym: &ast::PredicateSym) -> Option<usize> {
104        self.strata.iter().position(|x| x.contains(sym))
105    }
106}
107
108#[cfg(test)]
109mod test {
110    use super::*;
111    use std::collections::HashSet;
112
113    static FOO: ast::PredicateSym = ast::PredicateSym {
114        name: "foo",
115        arity: Some(2),
116    };
117    static BAR: ast::PredicateSym = ast::PredicateSym {
118        name: "bar",
119        arity: Some(1),
120    };
121
122    #[test]
123    fn try_eval() {
124        let bump = Bump::new();
125        let mut simple = SimpleProgram {
126            bump: &bump,
127            ext_preds: vec![FOO],
128            rules: HashMap::new(),
129        };
130
131        // Add a clause.
132        let atom = ast::Atom {
133            sym: FOO,
134            args: &[&ast::BaseTerm::Variable("X"), &ast::BaseTerm::Variable("_")],
135        };
136        let clause = &ast::Clause {
137            head: &ast::Atom {
138                sym: BAR,
139                args: &[&ast::BaseTerm::Variable("X")],
140            },
141            premises: &[&ast::Term::Atom(&atom)],
142            transform: &[],
143        };
144        simple.add_clause(clause);
145
146        assert_eq!(simple.extensional_preds().collect::<Vec<_>>(), vec![&FOO]);
147        assert_eq!(simple.intensional_preds().collect::<Vec<_>>(), vec![&BAR]);
148
149        let mut single_layer = HashSet::new();
150        single_layer.insert(&BAR);
151        let strata = vec![single_layer.clone()];
152        let stratified = simple.stratify(strata.iter());
153
154        assert_eq!(stratified.pred_to_index(&BAR), Some(0));
155        assert_eq!(stratified.strata().collect::<Vec<_>>(), vec![&single_layer])
156    }
157}
mangle_analysis/simple_program.rs

mangle_analysis/
simple_program.rs
