kodept_inference/
assumption.rs

1use std::borrow::{Borrow, Cow};
2use std::collections::HashMap;
3use std::convert::Infallible;
4use std::fmt::{Debug, Display, Formatter};
5use std::hash::Hash;
6use std::ops::Add;
7
8use itertools::Itertools;
9
10use crate::language::Var;
11use crate::r#type::{MonomorphicType, PolymorphicType};
12use crate::substitution::Substitutions;
13use crate::traits::{EnvironmentProvider, Substitutable};
14
15#[derive(PartialEq, Clone)]
16#[repr(transparent)]
17pub struct TypeTable<T, Name = Var>(HashMap<Name, T>)
18where
19    Name: Hash + Eq;
20
21pub type AssumptionSet = TypeTable<Vec<MonomorphicType>, Var>;
22pub type Environment = TypeTable<PolymorphicType, Var>;
23
24pub trait TypeTableOps {
25    fn merge(&mut self, other: Self);
26}
27
28// -------------------------------------------------------------------------------------------------
29
30impl<K: Hash + Eq, V> Default for TypeTable<V, K> {
31    fn default() -> Self {
32        Self(HashMap::new())
33    }
34}
35
36impl<K: Hash + Eq, V> TypeTable<V, K> {
37    pub fn empty() -> Self {
38        Self(HashMap::new())
39    }
40
41    pub fn remove<Q>(&mut self, key: &Q)
42    where
43        K: Borrow<Q>,
44        Q: Hash + Eq,
45    {
46        self.0.remove(key);
47    }
48
49    pub fn keys(&self) -> impl Iterator<Item = &K> {
50        self.0.keys()
51    }
52}
53
54impl<K, V> Add for TypeTable<V, K>
55where
56    K: Hash + Eq,
57    Self: TypeTableOps,
58{
59    type Output = Self;
60
61    fn add(mut self, rhs: Self) -> Self::Output {
62        self.merge(rhs);
63        self
64    }
65}
66
67impl<K, V> Add<&TypeTable<V, K>> for TypeTable<V, K>
68where
69    K: Hash + Eq,
70    Self: Clone + TypeTableOps,
71{
72    type Output = Self;
73
74    fn add(mut self, rhs: &TypeTable<V, K>) -> Self::Output {
75        self.merge(rhs.clone());
76        self
77    }
78}
79
80impl<K: Hash + Eq, V> Debug for TypeTable<V, K>
81where
82    Self: Display,
83{
84    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
85        write!(f, "{self}")
86    }
87}
88
89// -------------------------------------------------------------------------------------------------
90
91impl AssumptionSet {
92    pub fn push(&mut self, key: impl Into<Var>, value: impl Into<MonomorphicType>) {
93        self.0.entry(key.into()).or_default().push(value.into())
94    }
95
96    pub fn get<K>(&self, key: &K) -> Cow<[MonomorphicType]>
97    where
98        Var: Borrow<K>,
99        K: Hash + Eq,
100    {
101        match self.0.get(key) {
102            None => Cow::Owned(vec![]),
103            Some(x) => Cow::Borrowed(x),
104        }
105    }
106
107    pub fn single(key: impl Into<Var>, value: impl Into<MonomorphicType>) -> Self {
108        Self(HashMap::from([(key.into(), vec![value.into()])]))
109    }
110
111    pub fn merge_many(iter: impl IntoIterator<Item = AssumptionSet>) -> AssumptionSet {
112        iter.into_iter()
113            .fold(AssumptionSet::empty(), AssumptionSet::add)
114    }
115}
116
117impl TypeTableOps for AssumptionSet {
118    fn merge(&mut self, other: Self) {
119        for (k, v) in other.0 {
120            self.0.entry(k).or_default().extend(v)
121        }
122    }
123}
124
125impl Display for AssumptionSet {
126    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
127        write!(
128            f,
129            "[{}]",
130            self.0
131                .iter()
132                .map(|(k, v)| format!("{k} :: [{}]", v.iter().join(", ")))
133                .join(", ")
134        )
135    }
136}
137
138// -------------------------------------------------------------------------------------------------
139
140impl Environment {
141    pub fn substitute_mut(&mut self, substitutions: &Substitutions) -> &mut Self {
142        for t in self.0.values_mut() {
143            *t = t.substitute(substitutions);
144        }
145        self
146    }
147
148    pub fn push(&mut self, key: impl Into<Var>, value: impl Into<PolymorphicType>) {
149        self.0.insert(key.into(), value.into());
150    }
151
152    pub fn iter(&self) -> impl Iterator<Item = (&Var, &PolymorphicType)> {
153        self.0.iter()
154    }
155}
156
157impl TypeTableOps for Environment {
158    fn merge(&mut self, other: Self) {
159        self.0.extend(other.0)
160    }
161}
162
163impl Display for Environment {
164    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
165        write!(
166            f,
167            "[{}]",
168            self.0.iter().map(|(k, v)| format!("{k} :: {v}")).join(", ")
169        )
170    }
171}
172
173impl EnvironmentProvider<Var> for Environment {
174    type Error = Infallible;
175
176    fn maybe_get(&self, key: &Var) -> Result<Option<Cow<PolymorphicType>>, Self::Error> {
177        Ok(self.0.get(key).map(Cow::Borrowed))
178    }
179}