rue_types/
extractors.rs

1use id_arena::Arena;
2use indexmap::{IndexSet, indexset};
3
4use crate::{AtomRestriction, FunctionType, Pair, Type, TypeId, substitute};
5
6#[derive(Debug, Clone)]
7pub enum Atoms {
8    Unrestricted,
9    Restricted(IndexSet<AtomRestriction>),
10}
11
12pub fn extract_atoms(arena: &mut Arena<Type>, id: TypeId, strict: bool) -> Option<Atoms> {
13    let id = substitute(arena, id);
14    extract_atoms_impl(arena, id, strict)
15}
16
17fn extract_atoms_impl(arena: &Arena<Type>, id: TypeId, strict: bool) -> Option<Atoms> {
18    match arena[id].clone() {
19        Type::Apply(_) => unreachable!(),
20        Type::Ref(id) => extract_atoms_impl(arena, id, strict),
21        Type::Unresolved => Some(Atoms::Unrestricted),
22        Type::Generic(_) | Type::Never | Type::Function(_) | Type::Pair(_) | Type::Any => None,
23        Type::Atom(atom) => atom
24            .restriction
25            .map_or(Some(Atoms::Unrestricted), |restriction| {
26                Some(Atoms::Restricted(indexset![restriction]))
27            }),
28        Type::Struct(ty) => extract_atoms_impl(arena, ty.inner, strict),
29        Type::Alias(alias) => extract_atoms_impl(arena, alias.inner, strict),
30        Type::Union(ty) => {
31            let mut result = None;
32
33            for ty in ty.types {
34                let inner = extract_atoms_impl(arena, ty, strict);
35
36                let inner = if strict {
37                    inner?
38                } else if let Some(inner) = inner {
39                    inner
40                } else {
41                    continue;
42                };
43
44                match (&result, &inner) {
45                    (None, _) | (Some(_), Atoms::Unrestricted) => result = Some(inner),
46                    (Some(Atoms::Unrestricted), _) => {}
47                    (Some(Atoms::Restricted(restrictions)), Atoms::Restricted(inner)) => {
48                        let mut restrictions = restrictions.clone();
49                        restrictions.extend(inner.clone());
50                        result = Some(Atoms::Restricted(restrictions));
51                    }
52                }
53            }
54
55            result
56        }
57    }
58}
59
60pub fn extract_pairs(arena: &mut Arena<Type>, id: TypeId, strict: bool) -> Vec<Pair> {
61    let id = substitute(arena, id);
62    extract_pairs_impl(arena, id, strict).unwrap_or_default()
63}
64
65fn extract_pairs_impl(arena: &Arena<Type>, id: TypeId, strict: bool) -> Option<Vec<Pair>> {
66    match arena[id].clone() {
67        Type::Apply(_) => unreachable!(),
68        Type::Ref(id) => extract_pairs_impl(arena, id, strict),
69        Type::Unresolved => Some(vec![]),
70        Type::Generic(_) | Type::Never | Type::Atom(_) | Type::Function(_) | Type::Any => None,
71        Type::Pair(pair) => Some(vec![pair]),
72        Type::Struct(ty) => extract_pairs_impl(arena, ty.inner, strict),
73        Type::Alias(alias) => extract_pairs_impl(arena, alias.inner, strict),
74        Type::Union(ty) => {
75            let mut pairs = Vec::new();
76
77            for ty in ty.types {
78                let inner = extract_pairs_impl(arena, ty, strict);
79
80                if strict {
81                    pairs.extend(inner?);
82                } else {
83                    pairs.extend(inner.unwrap_or_default());
84                }
85            }
86
87            Some(pairs)
88        }
89    }
90}
91
92pub fn extract_functions(arena: &mut Arena<Type>, id: TypeId) -> Vec<FunctionType> {
93    let id = substitute(arena, id);
94    extract_functions_impl(arena, id).unwrap_or_default()
95}
96
97fn extract_functions_impl(arena: &Arena<Type>, id: TypeId) -> Option<Vec<FunctionType>> {
98    match arena[id].clone() {
99        Type::Apply(_) => unreachable!(),
100        Type::Ref(id) => extract_functions_impl(arena, id),
101        Type::Unresolved | Type::Never => Some(vec![]),
102        Type::Generic(_) | Type::Atom(_) | Type::Pair(_) | Type::Any => None,
103        Type::Function(function) => Some(vec![function]),
104        Type::Struct(ty) => extract_functions_impl(arena, ty.inner),
105        Type::Alias(alias) => extract_functions_impl(arena, alias.inner),
106        Type::Union(ty) => {
107            let mut pairs = Vec::new();
108
109            for ty in ty.types {
110                pairs.extend(extract_functions_impl(arena, ty)?);
111            }
112
113            Some(pairs)
114        }
115    }
116}
rue_types/extractors.rs

rue_types/
extractors.rs
