calyx_frontend/
lib_sig.rs

1use crate::Primitive;
2use calyx_utils::Id;
3use itertools::Itertools;
4use linked_hash_map::LinkedHashMap;
5use std::path::PathBuf;
6
7#[derive(Debug)]
8/// Tracks the information for [Primitive]s defined in the program.
9pub enum PrimitiveInfo {
10    /// An extern block that defines multiple primitives
11    Extern {
12        path: PathBuf,
13        primitives: LinkedHashMap<Id, Primitive>,
14        is_source: bool,
15    },
16    /// An inline primitive
17    Inline {
18        primitive: Primitive,
19        is_source: bool,
20    },
21}
22impl PrimitiveInfo {
23    pub fn ext(
24        path: PathBuf,
25        primitives: LinkedHashMap<Id, Primitive>,
26    ) -> Self {
27        PrimitiveInfo::Extern {
28            path,
29            primitives,
30            is_source: false,
31        }
32    }
33
34    pub fn inline(primitive: Primitive) -> Self {
35        PrimitiveInfo::Inline {
36            primitive,
37            is_source: false,
38        }
39    }
40
41    /// Check if this primitive is a source primitive
42    pub fn is_source(&self) -> bool {
43        match self {
44            PrimitiveInfo::Extern { is_source, .. } => *is_source,
45            PrimitiveInfo::Inline { is_source, .. } => *is_source,
46        }
47    }
48
49    /// Mark this primitive as a source primitive
50    pub fn set_source(&mut self) {
51        match self {
52            PrimitiveInfo::Extern { is_source, .. } => *is_source = true,
53            PrimitiveInfo::Inline { is_source, .. } => *is_source = true,
54        }
55    }
56}
57
58/// A representation of all the primitive definitions found while parsing
59/// the root program.
60#[derive(Debug, Default)]
61pub struct LibrarySignatures {
62    /// The primitives defined in the current context.
63    prims: Vec<PrimitiveInfo>,
64}
65
66impl LibrarySignatures {
67    /// Add a new inline primitive to the context.
68    /// Panics if a primitive with the same name is already defined.
69    pub fn add_inline_primitive(
70        &mut self,
71        primitive: Primitive,
72    ) -> &mut PrimitiveInfo {
73        assert!(
74            primitive.body.is_some(),
75            "inline primitive must have a body"
76        );
77        let name = primitive.name;
78        if self.find_primitive(name).is_some() {
79            panic!("Primitive `{name}` is already defined in the context.");
80        }
81        let prim = PrimitiveInfo::inline(primitive);
82        self.prims.push(prim);
83        self.prims.last_mut().unwrap()
84    }
85
86    /// Add a new, non-inline primitive to the context.
87    /// Panics if a primitive with the same name is already defined.
88    /// Requires that the file path is absolute and canonical.
89    pub fn add_extern_primitive(
90        &mut self,
91        file: PathBuf,
92        primitive: Primitive,
93    ) {
94        assert!(
95            primitive.body.is_none(),
96            "non-inline primitive must not have a body"
97        );
98        let name = primitive.name;
99        if self.find_primitive(name).is_some() {
100            panic!("Primitive `{name}` is already defined in the context.");
101        }
102        let definined_ext = self.prims.iter_mut().find(|prim| match prim {
103            PrimitiveInfo::Extern { path, .. } => path == &file,
104            _ => false,
105        });
106        if let Some(PrimitiveInfo::Extern { primitives, .. }) = definined_ext {
107            primitives.insert(name, primitive);
108        } else {
109            let mut primitives = LinkedHashMap::new();
110            primitives.insert(name, primitive);
111            self.prims.push(PrimitiveInfo::ext(file, primitives));
112        }
113    }
114
115    pub(crate) fn add_extern(
116        &mut self,
117        file: PathBuf,
118        prims: Vec<Primitive>,
119    ) -> &mut PrimitiveInfo {
120        let definined_ext = self.prims.iter().any(|prim| match prim {
121            PrimitiveInfo::Extern { path, .. } => path == &file,
122            _ => false,
123        });
124        if definined_ext {
125            panic!(
126                "Extern block with file `{}` is already defined in the context",
127                file.display()
128            );
129        }
130
131        let ext = PrimitiveInfo::ext(
132            file,
133            prims.into_iter().map(|p| (p.name, p)).collect(),
134        );
135        self.prims.push(ext);
136        self.prims.last_mut().unwrap()
137    }
138
139    /// Return the [Primitive] associated with the given name if defined, otherwise return None.
140    pub fn find_primitive<S>(&self, name: S) -> Option<&Primitive>
141    where
142        S: Into<Id>,
143    {
144        let key = name.into();
145        self.prims.iter().find_map(|prim| match prim {
146            PrimitiveInfo::Extern { primitives, .. } => primitives.get(&key),
147            PrimitiveInfo::Inline { primitive, .. } => {
148                if primitive.name == key {
149                    Some(primitive)
150                } else {
151                    None
152                }
153            }
154        })
155    }
156
157    /// Return the [Primitive] associated to this Id.
158    pub fn get_primitive<S>(&self, name: S) -> &Primitive
159    where
160        S: Into<Id>,
161    {
162        let key = name.into();
163        self.find_primitive(key).unwrap_or_else(|| {
164            panic!("Primitive `{key}` is not defined in the context.")
165        })
166    }
167
168    /// Mark an inlined primitive as a part of the source.
169    /// This is useful when using file mode compilation and printing only the source primitives.
170    /// Panics if the primitive is not defined.
171    pub fn mark_inline_source(&mut self, name: Id) {
172        let Some(inlined) = self.prims.iter_mut().find(|prim| match prim {
173            PrimitiveInfo::Inline { primitive, .. } => primitive.name == name,
174            PrimitiveInfo::Extern { .. } => false,
175        }) else {
176            panic!("Primitive `{name}` is not defined in the context.");
177        };
178        inlined.set_source()
179    }
180
181    /// Marks an `import`ed extern block as a part of the source.
182    /// There is no way to mark an individual primitive as a part of the source since the entire file will be linked.
183    /// Panics if the file path is not defined
184    pub fn mark_extern_source(&mut self, path: PathBuf) {
185        let Some(ext_def) = self.prims.iter_mut().find(|prim| match prim {
186            PrimitiveInfo::Extern { path: p, .. } => p == &path,
187            PrimitiveInfo::Inline { .. } => false,
188        }) else {
189            panic!(
190                "extern file `{}` is not defined in the context",
191                path.to_string_lossy()
192            );
193        };
194        ext_def.set_source()
195    }
196
197    /// Return an iterator over all defined primitives.
198    pub fn signatures(&self) -> impl Iterator<Item = &Primitive> + '_ {
199        self.prims.iter().flat_map(|prim| match prim {
200            PrimitiveInfo::Extern { primitives, .. } => {
201                primitives.values().collect_vec()
202            }
203            PrimitiveInfo::Inline { primitive, .. } => vec![primitive],
204        })
205    }
206
207    /// Returns all the underlying primitive information.
208    /// If you want all the signatures, use [LibrarySignatures::signatures] instead.
209    pub fn prim_infos(&self) -> &Vec<PrimitiveInfo> {
210        &self.prims
211    }
212
213    /// Return the underyling inlined primitives and whether they are source defined
214    pub fn prim_inlines(
215        &self,
216    ) -> impl Iterator<Item = (&Primitive, bool)> + '_ {
217        self.prims.iter().flat_map(|prim| match prim {
218            PrimitiveInfo::Extern { .. } => None,
219            PrimitiveInfo::Inline {
220                primitive,
221                is_source,
222            } => Some((primitive, *is_source)),
223        })
224    }
225
226    /// Return the paths for the extern defining files along with whether they are source defined.
227    pub fn extern_paths(&self) -> Vec<&PathBuf> {
228        self.prims
229            .iter()
230            .filter_map(|p| match p {
231                PrimitiveInfo::Extern { path, .. } => Some(path),
232                PrimitiveInfo::Inline { .. } => None,
233            })
234            .collect_vec()
235    }
236}