cozo/fixed_rule/utilities/
constant.rs

1/*
2 * Copyright 2022, The Cozo Project Authors.
3 *
4 * This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
5 * If a copy of the MPL was not distributed with this file,
6 * You can obtain one at https://mozilla.org/MPL/2.0/.
7 */
8
9use std::collections::BTreeMap;
10
11use miette::{bail, ensure, Diagnostic, Result};
12use smartstring::{LazyCompact, SmartString};
13use thiserror::Error;
14
15use crate::data::expr::Expr;
16use crate::data::program::WrongFixedRuleOptionError;
17use crate::data::symb::Symbol;
18use crate::data::value::DataValue;
19use crate::fixed_rule::{FixedRule, FixedRulePayload};
20use crate::parse::SourceSpan;
21use crate::runtime::db::Poison;
22use crate::runtime::temp_store::RegularTempStore;
23
24pub(crate) struct Constant;
25
26impl FixedRule for Constant {
27    fn run(
28        &self,
29        payload: FixedRulePayload<'_, '_>,
30        out: &mut RegularTempStore,
31        _poison: Poison,
32    ) -> Result<()> {
33        let data = payload.expr_option("data", None).unwrap();
34        let data = data.get_const().unwrap().get_slice().unwrap();
35        for row in data {
36            let tuple = row.get_slice().unwrap().into();
37            out.put(tuple)
38        }
39        Ok(())
40    }
41
42    fn arity(
43        &self,
44        options: &BTreeMap<SmartString<LazyCompact>, Expr>,
45        rule_head: &[Symbol],
46        span: SourceSpan,
47    ) -> Result<usize> {
48        let data = options
49            .get("data")
50            .unwrap()
51            .get_const()
52            .unwrap()
53            .get_slice()
54            .unwrap();
55        Ok(if data.is_empty() {
56            match rule_head.len() {
57                0 => {
58                    #[derive(Error, Debug, Diagnostic)]
59                    #[error("Constant rule does not have data")]
60                    #[diagnostic(code(parser::empty_const_rule))]
61                    #[diagnostic(help(
62                        "If you insist on using this empty rule, explicitly give its head"
63                    ))]
64                    struct EmptyConstRuleError(#[label] SourceSpan);
65                    bail!(EmptyConstRuleError(span))
66                }
67                i => i,
68            }
69        } else {
70            data.first().unwrap().get_slice().unwrap().len()
71        })
72    }
73
74    fn init_options(
75        &self,
76        options: &mut BTreeMap<SmartString<LazyCompact>, Expr>,
77        span: SourceSpan,
78    ) -> Result<()> {
79        let data = options
80            .get("data")
81            .ok_or_else(|| WrongFixedRuleOptionError {
82                name: "data".to_string(),
83                span: Default::default(),
84                rule_name: "Constant".to_string(),
85                help: "a list of lists is required".to_string(),
86            })?;
87        let data = match data.clone().eval_to_const()? {
88            DataValue::List(l) => l,
89            _ => bail!(WrongFixedRuleOptionError {
90                name: "data".to_string(),
91                span: Default::default(),
92                rule_name: "Constant".to_string(),
93                help: "a list of lists is required".to_string(),
94            }),
95        };
96
97        let mut tuples = vec![];
98        let mut last_len = None;
99        for row in data {
100            match row {
101                DataValue::List(tuple) => {
102                    if let Some(l) = &last_len {
103                        #[derive(Error, Debug, Diagnostic)]
104                        #[error("Constant head must have the same arity as the data given")]
105                        #[diagnostic(code(parser::const_data_arity_mismatch))]
106                        #[diagnostic(help("First row length: {0}; the mismatch: {1:?}"))]
107                        struct ConstRuleRowArityMismatch(
108                            usize,
109                            Vec<DataValue>,
110                            #[label] SourceSpan,
111                        );
112
113                        ensure!(
114                            *l == tuple.len(),
115                            ConstRuleRowArityMismatch(*l, tuple, span)
116                        );
117                    };
118                    last_len = Some(tuple.len());
119                    tuples.push(DataValue::List(tuple));
120                }
121                row => {
122                    #[derive(Error, Debug, Diagnostic)]
123                    #[error("Bad row for constant rule: {0:?}")]
124                    #[diagnostic(code(parser::bad_row_for_const))]
125                    #[diagnostic(help(
126                        "The body of a constant rule should evaluate to a list of lists"
127                    ))]
128                    struct ConstRuleRowNotList(DataValue);
129
130                    bail!(ConstRuleRowNotList(row))
131                }
132            }
133        }
134
135        options.insert(
136            SmartString::from("data"),
137            Expr::Const {
138                val: DataValue::List(tuples),
139                span: Default::default(),
140            },
141        );
142
143        Ok(())
144    }
145}