cozo/fixed_rule/utilities/
constant.rs1use 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}