spade_hir/
lib.rs

1pub mod expression;
2pub mod param_util;
3pub mod pretty_debug;
4pub mod pretty_print;
5pub mod query;
6pub mod symbol_table;
7pub mod testutil;
8
9use std::collections::{BTreeMap, HashMap};
10use std::fmt::Formatter;
11
12pub use expression::{Argument, ArgumentKind, ArgumentList, ExprKind, Expression};
13use itertools::Itertools;
14use num::BigInt;
15use serde::{Deserialize, Serialize};
16use spade_common::id_tracker::{ExprID, ImplID};
17use spade_common::{
18    location_info::{Loc, WithLocation},
19    name::{Identifier, NameID, Path},
20    num_ext::InfallibleToBigInt,
21};
22use spade_diagnostics::Diagnostic;
23use spade_types::{meta_types::MetaType, PrimitiveType};
24
25/**
26  Representation of the language with most language constructs still present, with
27  more correctness guaranatees than the AST, such as types actually existing.
28*/
29
30#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
31pub struct Block {
32    pub statements: Vec<Loc<Statement>>,
33    pub result: Option<Loc<Expression>>,
34}
35impl WithLocation for Block {}
36
37#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
38pub struct PatternArgument {
39    pub target: Loc<Identifier>,
40    pub value: Loc<Pattern>,
41    pub kind: ArgumentKind,
42}
43impl WithLocation for PatternArgument {}
44
45#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
46pub enum PatternKind {
47    Integer(BigInt),
48    Bool(bool),
49    Name {
50        name: Loc<NameID>,
51        pre_declared: bool,
52    },
53    Tuple(Vec<Loc<Pattern>>),
54    Array(Vec<Loc<Pattern>>),
55    /// Instantiation of an entity. While the argument contains information about
56    /// argument names, for codegen purposes, the arguments must be ordered in
57    /// the target order. I.e. they should all act as positioanl arguments
58    Type(Loc<NameID>, Vec<PatternArgument>),
59}
60impl PatternKind {
61    pub fn name(name: Loc<NameID>) -> Self {
62        PatternKind::Name {
63            name,
64            pre_declared: false,
65        }
66    }
67
68    pub fn integer(val: i32) -> Self {
69        Self::Integer(val.to_bigint())
70    }
71}
72impl PatternKind {
73    pub fn with_id(self, id: ExprID) -> Pattern {
74        Pattern { id, kind: self }
75    }
76
77    pub fn idless(self) -> Pattern {
78        Pattern {
79            id: ExprID(0),
80            kind: self,
81        }
82    }
83}
84impl std::fmt::Display for PatternKind {
85    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
86        match self {
87            PatternKind::Integer(val) => write!(f, "{val}"),
88            PatternKind::Bool(val) => write!(f, "{val}"),
89            PatternKind::Name { name, .. } => write!(f, "{name}"),
90            PatternKind::Tuple(members) => {
91                write!(
92                    f,
93                    "({})",
94                    members.iter().map(|m| format!("{}", m.kind)).join(", ")
95                )
96            }
97            PatternKind::Array(members) => {
98                write!(
99                    f,
100                    "[{}]",
101                    members.iter().map(|m| format!("{}", m.kind)).join(", ")
102                )
103            }
104            PatternKind::Type(name, _) => write!(f, "{name}(..)"),
105        }
106    }
107}
108
109#[derive(Debug, Clone, Serialize, Deserialize)]
110pub struct Pattern {
111    // Unique ID of the pattern for use in type inference. Shared with expressions
112    // meaning there are no expression/pattern id collisions
113    pub id: ExprID,
114    pub kind: PatternKind,
115}
116impl WithLocation for Pattern {}
117
118impl Pattern {
119    pub fn get_names(&self) -> Vec<Loc<NameID>> {
120        match &self.kind {
121            PatternKind::Integer(_) => vec![],
122            PatternKind::Bool(_) => vec![],
123            PatternKind::Name {
124                name,
125                pre_declared: _,
126            } => vec![name.clone()],
127            PatternKind::Tuple(inner) => inner.iter().flat_map(|i| i.get_names()).collect(),
128            PatternKind::Type(_, args) => {
129                args.iter().flat_map(|arg| arg.value.get_names()).collect()
130            }
131            PatternKind::Array(inner) => inner.iter().flat_map(|i| i.get_names()).collect(),
132        }
133    }
134}
135
136impl PartialEq for Pattern {
137    fn eq(&self, other: &Self) -> bool {
138        self.kind == other.kind
139    }
140}
141
142#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
143pub struct WalTrace {
144    pub clk: Option<Loc<Expression>>,
145    pub rst: Option<Loc<Expression>>,
146}
147impl WithLocation for WalTrace {}
148
149#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
150pub struct Binding {
151    pub pattern: Loc<Pattern>,
152    pub ty: Option<Loc<TypeSpec>>,
153    pub value: Loc<Expression>,
154    // Specifies if a wal_trace mir node should be emitted for this struct. If this
155    // is present, the type is traceable
156    pub wal_trace: Option<Loc<WalTrace>>,
157}
158
159#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
160pub enum PipelineRegMarkerExtra {
161    Condition(Loc<Expression>),
162    Count {
163        count: Loc<TypeExpression>,
164        count_typeexpr_id: ExprID,
165    },
166}
167
168#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
169pub enum Statement {
170    Error,
171    Binding(Binding),
172    Expression(Loc<Expression>),
173    Register(Register),
174    Declaration(Vec<Loc<NameID>>),
175    PipelineRegMarker(Option<PipelineRegMarkerExtra>),
176    Label(Loc<NameID>),
177    Assert(Loc<Expression>),
178    Set {
179        target: Loc<Expression>,
180        value: Loc<Expression>,
181    },
182    WalSuffixed {
183        suffix: Identifier,
184        target: Loc<NameID>,
185    },
186}
187impl WithLocation for Statement {}
188
189impl Statement {
190    /// NOTE: For use in tests
191    pub fn named_let(pattern_id: ExprID, name_id: Loc<NameID>, val: Expression) -> Self {
192        Self::Binding(Binding {
193            pattern: PatternKind::name(name_id).with_id(pattern_id).nowhere(),
194            ty: None,
195            value: val.nowhere(),
196            wal_trace: None,
197        })
198    }
199
200    pub fn binding(
201        pattern: Loc<Pattern>,
202        ty: Option<Loc<TypeSpec>>,
203        value: Loc<Expression>,
204    ) -> Statement {
205        Statement::Binding(Binding {
206            pattern,
207            ty,
208            value,
209            wal_trace: None,
210        })
211    }
212}
213
214#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
215pub struct Register {
216    pub pattern: Loc<Pattern>,
217    pub clock: Loc<Expression>,
218    pub reset: Option<(Loc<Expression>, Loc<Expression>)>,
219    pub initial: Option<Loc<Expression>>,
220    pub value: Loc<Expression>,
221    pub value_type: Option<Loc<TypeSpec>>,
222    pub attributes: AttributeList,
223}
224impl WithLocation for Register {}
225
226#[derive(PartialEq, Debug, Clone, PartialOrd, Eq, Ord, Serialize, Deserialize)]
227pub struct Module {
228    pub name: Loc<NameID>,
229    pub documentation: String,
230}
231
232/// Type params have both an identifier and a NameID since they go through the
233/// ast lowering process in a few separate steps, and the identifier needs to be
234/// re-added to the symtab multiple times
235#[derive(PartialEq, Debug, Clone, Hash, Eq, Serialize, Deserialize)]
236pub struct TypeParam {
237    pub ident: Loc<Identifier>,
238    pub name_id: NameID,
239    pub trait_bounds: Vec<Loc<TraitSpec>>,
240    pub meta: MetaType,
241}
242impl WithLocation for TypeParam {}
243impl TypeParam {
244    pub fn name_id(&self) -> NameID {
245        self.name_id.clone()
246    }
247}
248
249#[derive(PartialEq, Debug, Clone, Serialize, Deserialize, Hash, Eq)]
250pub enum TypeExpression {
251    /// An integer value
252    Integer(BigInt),
253    /// Another type
254    TypeSpec(TypeSpec),
255    ConstGeneric(Loc<ConstGeneric>),
256}
257impl WithLocation for TypeExpression {}
258
259impl TypeExpression {
260    fn replace_in(self, from: &TypeSpec, to: &TypeSpec) -> Self {
261        match self {
262            TypeExpression::TypeSpec(type_spec) => {
263                TypeExpression::TypeSpec(type_spec.replace_in(from, to))
264            }
265            TypeExpression::Integer(_) => self,
266            TypeExpression::ConstGeneric(_) => self,
267        }
268    }
269}
270
271impl std::fmt::Display for TypeExpression {
272    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
273        match self {
274            TypeExpression::Integer(val) => write!(f, "{val}"),
275            TypeExpression::TypeSpec(val) => write!(f, "{val}"),
276            TypeExpression::ConstGeneric(val) => write!(f, "{val}"),
277        }
278    }
279}
280
281/// A specification of a type to be used. For example, the types of input/output arguments the type
282/// of fields in a struct etc.
283#[derive(PartialEq, Debug, Clone, Serialize, Deserialize, Hash, Eq)]
284pub enum TypeSpec {
285    /// The type is a declared type (struct, enum, typedef etc.) with n arguments
286    Declared(Loc<NameID>, Vec<Loc<TypeExpression>>),
287    /// The type is a generic argument visible in the current scope
288    Generic(Loc<NameID>),
289    /// The type is a tuple of other variables
290    Tuple(Vec<Loc<TypeSpec>>),
291    Array {
292        inner: Box<Loc<TypeSpec>>,
293        size: Box<Loc<TypeExpression>>,
294    },
295    Inverted(Box<Loc<TypeSpec>>),
296    Wire(Box<Loc<TypeSpec>>),
297    /// The type of the `self` parameter in a trait method spec. Should not
298    /// occur in non-traits. The Loc is only used for diag_bails, so an approximate
299    /// reference is fine.
300    TraitSelf(Loc<()>),
301    /// A wildcard, cannot occur everywhere, but ast lowering enusres that wildcards are valid,
302    /// i.e. it is safe to emit a Diagnostic::bug if this is encountered where it is invalid
303    Wildcard(Loc<()>),
304}
305impl WithLocation for TypeSpec {}
306
307// Quick functions for creating types without typing so much
308impl TypeSpec {
309    pub fn unit() -> Self {
310        TypeSpec::Tuple(Vec::new())
311    }
312
313    pub fn type_params(&self) -> Vec<TypeExpression> {
314        match self {
315            TypeSpec::Declared(_, exprs) => exprs.clone().into_iter().map(|e| e.inner).collect(),
316            TypeSpec::Generic(_) => vec![],
317            TypeSpec::Tuple(_) => vec![],
318            TypeSpec::Array { inner, size } => {
319                vec![
320                    TypeExpression::TypeSpec(inner.inner.clone()),
321                    size.inner.clone(),
322                ]
323            }
324            TypeSpec::Inverted(inner) => vec![TypeExpression::TypeSpec(inner.inner.clone())],
325            TypeSpec::Wire(inner) => vec![TypeExpression::TypeSpec(inner.inner.clone())],
326            TypeSpec::TraitSelf(_) => vec![],
327            TypeSpec::Wildcard(_) => vec![],
328        }
329    }
330
331    pub fn replace_in(self, from: &TypeSpec, to: &TypeSpec) -> Self {
332        if &self == from {
333            to.clone()
334        } else {
335            match self {
336                TypeSpec::Declared(base, inner) => TypeSpec::Declared(
337                    base.clone(),
338                    inner
339                        .into_iter()
340                        .map(|expr| expr.map(|s| s.replace_in(from, to)))
341                        .collect(),
342                ),
343                TypeSpec::Generic(_) => self.clone(),
344                TypeSpec::Tuple(inner) => TypeSpec::Tuple(
345                    inner
346                        .into_iter()
347                        .map(|s| s.map(|s| s.replace_in(from, to)))
348                        .collect(),
349                ),
350                TypeSpec::Array { inner, size } => TypeSpec::Array {
351                    inner: Box::new(inner.map(|s| s.replace_in(from, to))),
352                    size: Box::new(size.map(|s| s.replace_in(from, to))),
353                },
354                TypeSpec::Inverted(inner) => {
355                    TypeSpec::Inverted(Box::new(inner.map(|s| s.replace_in(from, to))))
356                }
357                TypeSpec::Wire(inner) => {
358                    TypeSpec::Wire(Box::new(inner.map(|s| s.replace_in(from, to))))
359                }
360                TypeSpec::TraitSelf(_) => self,
361                TypeSpec::Wildcard(_) => self,
362            }
363        }
364    }
365}
366
367impl std::fmt::Display for TypeSpec {
368    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
369        let str = match self {
370            TypeSpec::Declared(name, params) => {
371                let type_params = if params.is_empty() {
372                    String::from("")
373                } else {
374                    format!(
375                        "<{}>",
376                        params
377                            .iter()
378                            .map(|g| format!("{g}"))
379                            .collect::<Vec<_>>()
380                            .join(", ")
381                    )
382                };
383                format!("{name}{type_params}")
384            }
385            TypeSpec::Generic(name) => format!("{name}"),
386            TypeSpec::Tuple(members) => {
387                format!(
388                    "({})",
389                    members
390                        .iter()
391                        .map(|m| format!("{m}"))
392                        .collect::<Vec<_>>()
393                        .join(", ")
394                )
395            }
396            TypeSpec::Array { inner, size } => format!("[{inner}; {size}]"),
397            TypeSpec::Inverted(inner) => format!("~{inner}"),
398            TypeSpec::Wire(inner) => format!("&{inner}"),
399            TypeSpec::TraitSelf(_) => "Self".into(),
400            TypeSpec::Wildcard(_) => "_".into(),
401        };
402        write!(f, "{str}")
403    }
404}
405
406/// A specification of a trait with type parameters
407#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Hash, Eq)]
408pub struct TraitSpec {
409    pub name: TraitName,
410    pub type_params: Option<Loc<Vec<Loc<TypeExpression>>>>,
411}
412impl WithLocation for TraitSpec {}
413
414/// Declaration of an enum
415#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
416pub struct Enum {
417    pub options: Vec<(Loc<NameID>, Loc<ParameterList>)>,
418    pub documentation: String,
419}
420impl WithLocation for Enum {}
421
422#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
423pub struct WalTraceable {
424    pub suffix: Path,
425    pub uses_clk: bool,
426    pub uses_rst: bool,
427}
428impl WithLocation for WalTraceable {}
429
430#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
431pub struct Struct {
432    pub members: Loc<ParameterList>,
433    pub is_port: bool,
434    pub attributes: AttributeList,
435    pub wal_traceable: Option<Loc<WalTraceable>>,
436    pub documentation: String,
437}
438impl WithLocation for Struct {}
439
440#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
441pub enum TypeDeclKind {
442    Enum(Loc<Enum>),
443    Primitive(PrimitiveType),
444    Struct(Loc<Struct>),
445}
446impl TypeDeclKind {
447    pub fn name(&self) -> &str {
448        match self {
449            TypeDeclKind::Enum(_) => "enum",
450            TypeDeclKind::Primitive(_) => "primitive",
451            TypeDeclKind::Struct(_) => "struct",
452        }
453    }
454}
455
456/// A declaration of a new type
457#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
458pub struct TypeDeclaration {
459    pub name: Loc<NameID>,
460    pub kind: TypeDeclKind,
461    pub generic_args: Vec<Loc<TypeParam>>,
462}
463impl WithLocation for TypeDeclaration {}
464
465#[derive(PartialEq, Debug, Clone, Serialize, Deserialize, Hash, Eq)]
466pub enum ConstGeneric {
467    Name(Loc<NameID>),
468    Const(BigInt),
469    Add(Box<Loc<ConstGeneric>>, Box<Loc<ConstGeneric>>),
470    Sub(Box<Loc<ConstGeneric>>, Box<Loc<ConstGeneric>>),
471    Mul(Box<Loc<ConstGeneric>>, Box<Loc<ConstGeneric>>),
472    Div(Box<Loc<ConstGeneric>>, Box<Loc<ConstGeneric>>),
473    Mod(Box<Loc<ConstGeneric>>, Box<Loc<ConstGeneric>>),
474    UintBitsToFit(Box<Loc<ConstGeneric>>),
475    Eq(Box<Loc<ConstGeneric>>, Box<Loc<ConstGeneric>>),
476    NotEq(Box<Loc<ConstGeneric>>, Box<Loc<ConstGeneric>>),
477}
478impl WithLocation for ConstGeneric {}
479
480impl ConstGeneric {
481    pub fn with_id(self, id: ExprID) -> ConstGenericWithId {
482        ConstGenericWithId { id, inner: self }
483    }
484}
485
486#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
487pub struct ConstGenericWithId {
488    pub id: ExprID,
489    pub inner: ConstGeneric,
490}
491impl WithLocation for ConstGenericWithId {}
492
493impl std::fmt::Display for ConstGeneric {
494    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
495        match self {
496            ConstGeneric::Name(n) => write!(f, "{n}"),
497            ConstGeneric::Const(val) => write!(f, "{val}"),
498            ConstGeneric::Add(l, r) => write!(f, "({l} + {r})"),
499            ConstGeneric::Sub(l, r) => write!(f, "({l} - {r})"),
500            ConstGeneric::Mul(l, r) => write!(f, "({l} * {r})"),
501            ConstGeneric::Div(l, r) => write!(f, "({l} / {r})"),
502            ConstGeneric::Mod(l, r) => write!(f, "({l} % {r})"),
503            ConstGeneric::Eq(l, r) => write!(f, "({l} == {r})"),
504            ConstGeneric::NotEq(l, r) => write!(f, "({l} != {r})"),
505            ConstGeneric::UintBitsToFit(a) => write!(f, "uint_bits_to_fit({a})"),
506        }
507    }
508}
509
510#[derive(PartialEq, Debug, Clone, Serialize, Deserialize, Hash, Eq)]
511pub enum WhereClause {
512    Int {
513        target: Loc<NameID>,
514        constraint: Loc<ConstGeneric>,
515    },
516    Type {
517        target: Loc<NameID>,
518        traits: Vec<Loc<TraitSpec>>,
519    },
520}
521impl WithLocation for WhereClause {}
522
523impl std::fmt::Display for WhereClause {
524    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
525        let str = match self {
526            WhereClause::Int { target, constraint } => {
527                format!("{target}: {{ {constraint} }}")
528            }
529            WhereClause::Type { target, traits } => {
530                format!(
531                    "{target}: {}",
532                    traits.iter().map(|trait_spec| &trait_spec.name).join(" + ")
533                )
534            }
535        };
536        write!(f, "{}", str)
537    }
538}
539
540#[derive(PartialEq, Debug, Clone, Serialize, Deserialize, Hash, Eq)]
541pub enum UnitName {
542    /// The name will be mangled down to contain the NameID in order to ensure
543    /// uniqueness. Emitted by generic functions
544    WithID(Loc<NameID>),
545    /// The name will contain the full path to the name but the ID section of the
546    /// nameID will not be included. Used by non-generic functions
547    FullPath(Loc<NameID>),
548    /// The name will not be mangled. In the output code it will appear as String
549    /// but the compiler will still refer to it by the NameID
550    Unmangled(String, Loc<NameID>),
551}
552
553impl UnitName {
554    pub fn name_id(&self) -> &Loc<NameID> {
555        match self {
556            UnitName::WithID(name) => name,
557            UnitName::FullPath(name) => name,
558            UnitName::Unmangled(_, name) => name,
559        }
560    }
561}
562
563impl std::fmt::Display for UnitName {
564    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
565        match self {
566            UnitName::WithID(name) | UnitName::FullPath(name) | UnitName::Unmangled(_, name) => {
567                write!(f, "{name}")
568            }
569        }
570    }
571}
572
573#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
574pub struct Unit {
575    pub name: UnitName,
576    pub head: UnitHead,
577    pub attributes: AttributeList,
578    // This is needed here because the head does not have NameIDs
579    pub inputs: Vec<(Loc<NameID>, Loc<TypeSpec>)>,
580    pub body: Loc<Expression>,
581}
582impl WithLocation for Unit {}
583
584#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
585pub struct Parameter {
586    /// If the #[no_mangle] attribute is present, this field is set
587    /// with the Loc pointing to the attribute
588    pub no_mangle: Option<Loc<()>>,
589    pub name: Loc<Identifier>,
590    pub ty: Loc<TypeSpec>,
591}
592
593#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
594pub struct ParameterList(pub Vec<Parameter>);
595impl WithLocation for ParameterList {}
596
597impl ParameterList {
598    pub fn argument_num(&self) -> usize {
599        self.0.len()
600    }
601
602    /// Look up the type of an argument. Panics if no such argument exists
603    pub fn arg_type(&self, name: &Identifier) -> &TypeSpec {
604        if let Some(result) = self.try_get_arg_type(name) {
605            result
606        } else {
607            panic!(
608                "Tried to get type of an argument which is not part of the parameter list. {}",
609                name
610            )
611        }
612    }
613
614    /// Look up the type of an argument, returning None if no such argument exists
615    pub fn try_get_arg_type(&self, name: &Identifier) -> Option<&Loc<TypeSpec>> {
616        for Parameter {
617            name: arg,
618            ty,
619            no_mangle: _,
620        } in &self.0
621        {
622            if &arg.inner == name {
623                return Some(ty);
624            }
625        }
626        None
627    }
628
629    pub fn arg_index(&self, target: &Identifier) -> Option<usize> {
630        let indices = self
631            .0
632            .iter()
633            .enumerate()
634            .filter_map(
635                |(
636                    i,
637                    Parameter {
638                        name,
639                        ty: _,
640                        no_mangle: _,
641                    },
642                )| {
643                    if &name.inner == target {
644                        Some(i)
645                    } else {
646                        None
647                    }
648                },
649            )
650            .collect::<Vec<_>>();
651
652        if indices.len() > 1 {
653            panic!("Duplicate arguments with the same name")
654        } else {
655            indices.first().cloned()
656        }
657    }
658}
659
660#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
661pub enum FunctionKind {
662    Fn,
663    Struct,
664    Enum,
665}
666
667#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
668pub enum UnitKind {
669    Function(FunctionKind),
670    Entity,
671    Pipeline {
672        depth: Loc<TypeExpression>,
673        depth_typeexpr_id: ExprID,
674    },
675}
676impl WithLocation for UnitKind {}
677
678impl UnitKind {
679    pub fn name(&self) -> &'static str {
680        match self {
681            UnitKind::Function(FunctionKind::Fn) => "function",
682            UnitKind::Function(FunctionKind::Struct) => "struct",
683            UnitKind::Function(FunctionKind::Enum) => "enum variant",
684            UnitKind::Entity => "entity",
685            UnitKind::Pipeline { .. } => "pipeline",
686        }
687    }
688
689    pub fn is_pipeline(&self) -> bool {
690        matches!(self, UnitKind::Pipeline { .. })
691    }
692}
693
694#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
695pub struct UnitHead {
696    pub name: Loc<Identifier>,
697    pub inputs: Loc<ParameterList>,
698    /// (-> token, type)
699    pub output_type: Option<Loc<TypeSpec>>,
700    pub unit_type_params: Vec<Loc<TypeParam>>,
701    pub scope_type_params: Vec<Loc<TypeParam>>,
702    pub unit_kind: Loc<UnitKind>,
703    pub where_clauses: Vec<Loc<WhereClause>>,
704    pub documentation: String,
705}
706impl WithLocation for UnitHead {}
707
708impl UnitHead {
709    pub fn output_type(&self) -> Loc<TypeSpec> {
710        match &self.output_type {
711            Some(t) => t.clone(),
712            None => {
713                // FIXME: We should point to the end of the argument list here
714                TypeSpec::unit().at_loc(&self.name.loc())
715            }
716        }
717    }
718    pub fn get_type_params(&self) -> Vec<Loc<TypeParam>> {
719        self.unit_type_params
720            .iter()
721            .chain(self.scope_type_params.iter())
722            .cloned()
723            .collect_vec()
724    }
725}
726
727#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
728pub enum Item {
729    Unit(Loc<Unit>),
730    ExternUnit(UnitName, Loc<UnitHead>),
731}
732
733impl Item {
734    pub fn assume_unit(&self) -> &Unit {
735        match self {
736            Item::Unit(u) => &u.inner,
737            Item::ExternUnit(_, _) => panic!("Expected unit, got extern unit"),
738        }
739    }
740}
741
742/// Items which have associated code that can be executed. This is different from
743/// type declarations which are items, but which do not have code on their own
744#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
745pub enum ExecutableItem {
746    EnumInstance { base_enum: NameID, variant: usize },
747    StructInstance,
748    Unit(Loc<Unit>),
749    ExternUnit(UnitName, Loc<UnitHead>),
750}
751impl WithLocation for ExecutableItem {}
752
753pub type TypeList = HashMap<NameID, Loc<TypeDeclaration>>;
754
755#[derive(Serialize, Deserialize, PartialEq, Eq, Hash, Debug, Clone, PartialOrd, Ord)]
756pub enum TraitName {
757    Named(Loc<NameID>),
758    Anonymous(ImplID),
759}
760
761impl TraitName {
762    pub fn is_anonymous(&self) -> bool {
763        matches!(self, Self::Anonymous(_))
764    }
765
766    /// Returns the loc of the name of this trait, if it exists. None otherwise
767    pub fn name_loc(&self) -> Option<Loc<NameID>> {
768        match self {
769            TraitName::Named(n) => Some(n.clone()),
770            TraitName::Anonymous(_) => None,
771        }
772    }
773}
774
775impl std::fmt::Display for TraitName {
776    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
777        match self {
778            TraitName::Named(n) => write!(f, "{n}"),
779            TraitName::Anonymous(id) => write!(f, "Anonymous({})", id.0),
780        }
781    }
782}
783
784/// Attributes that are still present as attributes in the HIR. Not all AST
785/// attributes are in this list, as some are consumed and inlined into the corresponding
786/// ast node
787#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
788pub enum Attribute {
789    Optimize { passes: Vec<Loc<String>> },
790    Fsm { state: NameID },
791    WalTraceable { suffix: Identifier },
792}
793impl Attribute {
794    pub fn name(&self) -> &str {
795        match self {
796            Attribute::Optimize { passes: _ } => "optimize",
797            Attribute::Fsm { state: _ } => "fsm",
798            Attribute::WalTraceable { suffix: _ } => "suffix",
799        }
800    }
801}
802impl WithLocation for Attribute {}
803
804#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
805pub struct AttributeList(pub Vec<Loc<Attribute>>);
806
807impl AttributeList {
808    pub fn empty() -> Self {
809        Self(vec![])
810    }
811}
812
813#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
814pub struct ImplBlock {
815    /// Mapping of identifiers to the NameID of the entity which is the implementation
816    /// for the specified function
817    pub fns: HashMap<Identifier, (NameID, Loc<()>)>,
818    pub type_params: Vec<Loc<TypeParam>>,
819    pub target: Loc<TypeSpec>,
820    pub id: ImplID,
821}
822impl WithLocation for ImplBlock {}
823
824#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
825pub struct TraitDef {
826    pub type_params: Option<Loc<Vec<Loc<TypeParam>>>>,
827    pub fns: HashMap<Identifier, Loc<UnitHead>>,
828}
829impl WithLocation for TraitDef {}
830
831#[derive(PartialEq, Hash, Eq, Debug, Clone, Serialize, Deserialize)]
832pub enum ImplTarget {
833    Array,
834    Inverted,
835    Wire,
836    Named(NameID),
837}
838
839impl ImplTarget {
840    pub fn display(&self, args: &[TypeExpression]) -> String {
841        match self {
842            ImplTarget::Array => {
843                format!(
844                    "[{}; {}]",
845                    args.get(0)
846                        .map(|a| format!("{}", a))
847                        .unwrap_or_else(|| "<(bug) Missing param 0>".to_string()),
848                    args.get(1)
849                        .map(|a| format!("{}", a))
850                        .unwrap_or_else(|| "<(bug) Missing param 1>".to_string())
851                )
852            }
853            ImplTarget::Wire => {
854                format!(
855                    "&{}",
856                    args.get(0)
857                        .map(|a| format!("{}", a))
858                        .unwrap_or_else(|| "<(bug) Missing param 0>".to_string()),
859                )
860            }
861            ImplTarget::Inverted => {
862                format!(
863                    "inv {}",
864                    args.get(0)
865                        .map(|a| format!("{}", a))
866                        .unwrap_or_else(|| "<(bug) Missing param 0>".to_string()),
867                )
868            }
869            ImplTarget::Named(name) => {
870                format!(
871                    "{}{}",
872                    name,
873                    if args.is_empty() {
874                        format!("")
875                    } else {
876                        format!("<{}>", args.iter().map(|arg| format!("{}", arg)).join(", "))
877                    }
878                )
879            }
880        }
881    }
882}
883
884/// A list of all the items present in the whole AST, flattened to remove module
885/// hierarchies.
886///
887/// That is, `mod a { mod b{ entity X {} } } will result in members containing `a::b::X`, but the
888/// module hierarchy no longer has to be traversed.
889///
890/// The modules themselves however still exist in the `modules` set.
891#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
892pub struct ItemList {
893    pub executables: BTreeMap<NameID, ExecutableItem>,
894    pub types: TypeList,
895    /// All modules, including empty ones.
896    pub modules: BTreeMap<NameID, Module>,
897    // FIXME: Support entities and pipelines as trait members.
898    /// All traits in the compilation unit. Traits consist of a list of functions
899    /// by name. Anonymous impl blocks are also members here, but their name is never
900    /// visible to the user.
901    pub traits: HashMap<TraitName, TraitDef>,
902    pub impls: HashMap<ImplTarget, HashMap<(TraitName, Vec<TypeExpression>), Loc<ImplBlock>>>,
903}
904
905impl Default for ItemList {
906    fn default() -> Self {
907        Self::new()
908    }
909}
910
911impl ItemList {
912    pub fn new() -> Self {
913        Self {
914            executables: BTreeMap::new(),
915            types: TypeList::new(),
916            modules: BTreeMap::new(),
917            traits: HashMap::new(),
918            impls: HashMap::new(),
919        }
920    }
921
922    pub fn add_executable(
923        &mut self,
924        name: Loc<NameID>,
925        item: ExecutableItem,
926    ) -> Result<(), Diagnostic> {
927        if let Some(_) = self.executables.get_key_value(&name) {
928            Err(
929                Diagnostic::error(&name, format!("Multiple definitions of thing {name}"))
930                    .primary_label("New definition"),
931            )
932        } else {
933            self.executables.insert(name.inner, item);
934            Ok(())
935        }
936    }
937
938    pub fn add_trait(
939        &mut self,
940        name: TraitName,
941        type_params: Option<Loc<Vec<Loc<TypeParam>>>>,
942        members: Vec<(Identifier, Loc<UnitHead>)>,
943    ) -> Result<(), Diagnostic> {
944        if let Some((prev, _)) = self.traits.get_key_value(&name) {
945            Err(
946                // NOTE: unwrap here is safe *unless* we have got duplicate trait IDs which
947                // means we have bigger problems
948                Diagnostic::error(
949                    name.name_loc().unwrap(),
950                    format!("Multiple definitions of trait {name}"),
951                )
952                .primary_label("New definition")
953                .secondary_label(prev.name_loc().unwrap(), "Previous definition"),
954            )
955        } else {
956            self.traits.insert(
957                name,
958                TraitDef {
959                    type_params,
960                    fns: members.into_iter().collect(),
961                },
962            );
963            Ok(())
964        }
965    }
966
967    pub fn get_trait(&self, name: &TraitName) -> Option<&TraitDef> {
968        self.traits.get(name)
969    }
970
971    pub fn traits(&self) -> &HashMap<TraitName, TraitDef> {
972        &self.traits
973    }
974}