arithmetic_typing/env/
mod.rs

1//! `TypeEnvironment` and related types.
2
3use std::{collections::HashMap, iter::FromIterator, ops};
4
5use crate::{
6    arith::{
7        Constraint, ConstraintSet, MapPrimitiveType, Num, NumArithmetic, ObjectSafeConstraint,
8        Substitutions, TypeArithmetic,
9    },
10    ast::TypeAst,
11    error::Errors,
12    types::{ParamConstraints, ParamQuantifier},
13    PrimitiveType, Type,
14};
15use arithmetic_parser::{grammars::Grammar, Block};
16
17mod processor;
18
19use self::processor::TypeProcessor;
20
21/// Environment containing type information on named variables.
22///
23/// # Examples
24///
25/// See [the crate docs](index.html#examples) for examples of usage.
26///
27/// # Concrete and partially specified types
28///
29/// The environment retains full info on the types even if the type is not
30/// [concrete](Type::is_concrete()). Non-concrete types are tied to an environment.
31/// An environment will panic on inserting a non-concrete type via [`Self::insert()`]
32/// or other methods.
33///
34/// ```
35/// # use arithmetic_parser::grammars::{F32Grammar, Parse};
36/// # use arithmetic_typing::{defs::Prelude, Annotated, TypeEnvironment};
37/// # type Parser = Annotated<F32Grammar>;
38/// # fn main() -> anyhow::Result<()> {
39/// // An easy way to get a non-concrete type is to involve `any`.
40/// let code = "(x, ...) = (1, 2, 3) as any;";
41/// let code = Parser::parse_statements(code)?;
42///
43/// let mut env: TypeEnvironment = Prelude::iter().collect();
44/// env.process_statements(&code)?;
45/// assert!(!env["x"].is_concrete());
46/// # Ok(())
47/// # }
48/// ```
49#[derive(Debug, Clone)]
50pub struct TypeEnvironment<Prim: PrimitiveType = Num> {
51    pub(crate) substitutions: Substitutions<Prim>,
52    pub(crate) known_constraints: ConstraintSet<Prim>,
53    variables: HashMap<String, Type<Prim>>,
54}
55
56impl<Prim: PrimitiveType> Default for TypeEnvironment<Prim> {
57    fn default() -> Self {
58        Self {
59            variables: HashMap::new(),
60            known_constraints: Prim::well_known_constraints(),
61            substitutions: Substitutions::default(),
62        }
63    }
64}
65
66impl<Prim: PrimitiveType> TypeEnvironment<Prim> {
67    /// Creates an empty environment.
68    pub fn new() -> Self {
69        Self::default()
70    }
71
72    /// Gets type of the specified variable.
73    pub fn get(&self, name: &str) -> Option<&Type<Prim>> {
74        self.variables.get(name)
75    }
76
77    /// Iterates over variables contained in this env.
78    pub fn iter(&self) -> impl Iterator<Item = (&str, &Type<Prim>)> + '_ {
79        self.variables.iter().map(|(name, ty)| (name.as_str(), ty))
80    }
81
82    fn prepare_type(ty: impl Into<Type<Prim>>) -> Type<Prim> {
83        let mut ty = ty.into();
84        assert!(ty.is_concrete(), "Type {} is not concrete", ty);
85
86        if let Type::Function(function) = &mut ty {
87            if function.params.is_none() {
88                ParamQuantifier::set_params(function, ParamConstraints::default());
89            }
90        }
91        ty
92    }
93
94    /// Sets type of a variable.
95    ///
96    /// # Panics
97    ///
98    /// - Will panic if `ty` is not [concrete](Type::is_concrete()). Non-concrete
99    ///   types are tied to the environment; inserting them into an env is a logical error.
100    pub fn insert(&mut self, name: &str, ty: impl Into<Type<Prim>>) -> &mut Self {
101        self.variables
102            .insert(name.to_owned(), Self::prepare_type(ty));
103        self
104    }
105
106    /// Inserts a [`Constraint`] into the environment so that it can be used when parsing
107    /// type annotations.
108    ///
109    /// Adding a constraint is not mandatory for it to be usable during type inference;
110    /// this method only influences whether the constraint is recognized during type parsing.
111    pub fn insert_constraint(&mut self, constraint: impl Constraint<Prim>) -> &mut Self {
112        self.known_constraints.insert(constraint);
113        self
114    }
115
116    /// Inserts an [`ObjectSafeConstraint`] into the environment so that it can be used
117    /// when parsing type annotations.
118    ///
119    /// Other than more strict type requirements, this method is identical to
120    /// [`Self::insert_constraint`].
121    pub fn insert_object_safe_constraint(
122        &mut self,
123        constraint: impl ObjectSafeConstraint<Prim>,
124    ) -> &mut Self {
125        self.known_constraints.insert_object_safe(constraint);
126        self
127    }
128
129    /// Processes statements with the default type arithmetic. After processing, the environment
130    /// will contain type info about newly declared vars.
131    ///
132    /// This method is a shortcut for calling `process_with_arithmetic` with
133    /// [`NumArithmetic::without_comparisons()`].
134    pub fn process_statements<'a, T>(
135        &mut self,
136        block: &Block<'a, T>,
137    ) -> Result<Type<Prim>, Errors<'a, Prim>>
138    where
139        T: Grammar<'a, Type = TypeAst<'a>>,
140        NumArithmetic: MapPrimitiveType<T::Lit, Prim = Prim> + TypeArithmetic<Prim>,
141    {
142        self.process_with_arithmetic(&NumArithmetic::without_comparisons(), block)
143    }
144
145    /// Processes statements with a given `arithmetic`. After processing, the environment
146    /// will contain type info about newly declared vars.
147    ///
148    /// # Errors
149    ///
150    /// Even if there are any type errors, all statements in the `block` will be executed
151    /// to completion and all errors will be reported. However, the environment will **not**
152    /// include any vars beyond the first failing statement.
153    pub fn process_with_arithmetic<'a, T, A>(
154        &mut self,
155        arithmetic: &A,
156        block: &Block<'a, T>,
157    ) -> Result<Type<Prim>, Errors<'a, Prim>>
158    where
159        T: Grammar<'a, Type = TypeAst<'a>>,
160        A: MapPrimitiveType<T::Lit, Prim = Prim> + TypeArithmetic<Prim>,
161    {
162        TypeProcessor::new(self, arithmetic).process_statements(block)
163    }
164}
165
166impl<Prim: PrimitiveType> ops::Index<&str> for TypeEnvironment<Prim> {
167    type Output = Type<Prim>;
168
169    fn index(&self, name: &str) -> &Self::Output {
170        self.get(name)
171            .unwrap_or_else(|| panic!("Variable `{}` is not defined", name))
172    }
173}
174
175fn convert_iter<Prim: PrimitiveType, S, Ty, I>(
176    iter: I,
177) -> impl Iterator<Item = (String, Type<Prim>)>
178where
179    I: IntoIterator<Item = (S, Ty)>,
180    S: Into<String>,
181    Ty: Into<Type<Prim>>,
182{
183    iter.into_iter()
184        .map(|(name, ty)| (name.into(), TypeEnvironment::prepare_type(ty)))
185}
186
187impl<Prim: PrimitiveType, S, Ty> FromIterator<(S, Ty)> for TypeEnvironment<Prim>
188where
189    S: Into<String>,
190    Ty: Into<Type<Prim>>,
191{
192    fn from_iter<I: IntoIterator<Item = (S, Ty)>>(iter: I) -> Self {
193        Self {
194            variables: convert_iter(iter).collect(),
195            known_constraints: Prim::well_known_constraints(),
196            substitutions: Substitutions::default(),
197        }
198    }
199}
200
201impl<Prim: PrimitiveType, S, Ty> Extend<(S, Ty)> for TypeEnvironment<Prim>
202where
203    S: Into<String>,
204    Ty: Into<Type<Prim>>,
205{
206    fn extend<I: IntoIterator<Item = (S, Ty)>>(&mut self, iter: I) {
207        self.variables.extend(convert_iter(iter))
208    }
209}
210
211// Helper trait to wrap type mapper and arithmetic.
212trait FullArithmetic<Val, Prim: PrimitiveType>:
213    MapPrimitiveType<Val, Prim = Prim> + TypeArithmetic<Prim>
214{
215}
216
217impl<Val, Prim: PrimitiveType, T> FullArithmetic<Val, Prim> for T where
218    T: MapPrimitiveType<Val, Prim = Prim> + TypeArithmetic<Prim>
219{
220}
arithmetic_typing/env/mod.rs

arithmetic_typing/env/
mod.rs
