pliron_llvm/
types.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) The pliron contributors
3
4//! [Type]s defined in the LLVM dialect.
5
6use core::hash::Hash;
7use pliron::{
8    builtin::type_interfaces::FunctionTypeInterface,
9    combine::{Parser, between, optional, token},
10    common_traits::Verify,
11    context::Context,
12    derive::{format, pliron_type, type_interface_impl},
13    identifier::Identifier,
14    input_err_noloc,
15    irfmt::{
16        parsers::{delimited_list_parser, location, spaced, type_parser},
17        printers::{enclosed, list_with_sep},
18    },
19    location::Located,
20    parsable::{IntoParseResult, Parsable, ParseResult, StateStream},
21    printable::{self, ListSeparator, Printable},
22    result::Result,
23    r#type::{Type, TypeHandle, TypedHandle},
24    verify_err_noloc,
25};
26use thiserror::Error;
27
28/// Represents a c-like struct type.
29/// Limitations and warnings on its usage are similar to that in MLIR.
30/// `<https://mlir.llvm.org/docs/Dialects/LLVM/#structure-types>`
31///   1. Anonymous (aka unnamed) structs cannot be recursive.
32///   2. Named structs are uniqued *only* by name, and may be recursive.
33///   3. LLVM calls anonymous structs as literal structs and
34///      named structs as identified structs.
35///   4. Named structs may be opaque, i.e., no body specificed.
36///      Recursive types may be created by first creating an opaque struct
37///      and later setting its fields (body).
38#[pliron_type(name = "llvm.struct")]
39#[derive(Debug)]
40pub struct StructType {
41    name: Option<Identifier>,
42    fields: Option<Vec<TypeHandle>>,
43}
44
45impl StructType {
46    /// Get or create a named StructType.
47    /// If `fields` is `None`, it indicates an opaque struct.
48    /// A body can be added to opaque structs by calling this again later.
49    /// Returns an error if all of the below conditions are true:
50    ///   a. The name is already registered
51    ///   b. The body is already set (i.e, the struct is not oqaue)
52    ///   c. The fields provided here don't match with the existing body.
53    /// Since named structs only rely on the name for uniqueness,
54    /// It is not an error to provide `fields` as `None` even when
55    /// the named struct already exists and has its body set.
56    pub fn get_named(
57        ctx: &Context,
58        name: Identifier,
59        fields: Option<Vec<TypeHandle>>,
60    ) -> Result<TypedHandle<Self>> {
61        let self_handle = Type::instantiate(
62            StructType {
63                name: Some(name.clone()),
64                // Uniquing happens only on the name, so this doesn't matter.
65                fields: None,
66            },
67            ctx,
68        );
69        // Verify that we created a new or equivalent existing type.
70        let mut self_ref = self_handle.to_handle().deref_mut(ctx);
71        let self_ref = self_ref.downcast_mut::<StructType>().unwrap();
72        assert!(self_ref.name.as_ref().unwrap() == &name);
73        if let Some(fields) = fields {
74            // We've been provided fields to be set.
75            if let Some(existing_fields) = &self_ref.fields {
76                // Fields were already set before, ensure they're same as the given ones.
77                if existing_fields != &fields {
78                    input_err_noloc!(StructErr::ExistingMismatch(name.into()))?
79                }
80            } else {
81                // Set the fields now.
82                self_ref.fields = Some(fields);
83            }
84        }
85        Ok(self_handle)
86    }
87
88    /// Get or create a new unnamed (anonymous) struct.
89    /// These are finalized upon creation, and uniqued based on the fields.
90    pub fn get_unnamed(ctx: &Context, fields: Vec<TypeHandle>) -> TypedHandle<Self> {
91        Type::instantiate(
92            StructType {
93                name: None,
94                fields: Some(fields),
95            },
96            ctx,
97        )
98    }
99
100    /// Does this struct not have its body set?
101    pub fn is_opaque(&self) -> bool {
102        self.fields.is_none()
103    }
104
105    /// Is this a named struct?
106    pub fn is_named(&self) -> bool {
107        self.name.is_some()
108    }
109
110    /// Get this struct's name, if it has one.
111    pub fn name(&self) -> Option<Identifier> {
112        self.name.clone()
113    }
114
115    /// Get type of the idx'th field.
116    pub fn field_type(&self, field_idx: usize) -> TypeHandle {
117        self.fields
118            .as_ref()
119            .expect("field_type shouldn't be called on opaque types")[field_idx]
120    }
121
122    /// Get the number of fields this struct has
123    pub fn num_fields(&self) -> usize {
124        self.fields
125            .as_ref()
126            .expect("num_fields shouldn't be called on opaque types")
127            .len()
128    }
129
130    /// Get an iterator over the fields of this struct
131    pub fn fields(&self) -> impl Iterator<Item = TypeHandle> + '_ {
132        self.fields
133            .as_ref()
134            .expect("fields shouldn't be called on opaque types")
135            .iter()
136            .cloned()
137    }
138}
139
140#[derive(Debug, Error)]
141pub enum StructErr {
142    #[error("struct cannot be both opaque and anonymous")]
143    OpaqueAndAnonymousErr,
144    #[error("struct {0} already exists and is different")]
145    ExistingMismatch(String),
146}
147
148impl Verify for StructType {
149    fn verify(&self, _ctx: &Context) -> Result<()> {
150        if self.name.is_none() && self.fields.is_none() {
151            verify_err_noloc!(StructErr::OpaqueAndAnonymousErr)?
152        }
153        Ok(())
154    }
155}
156
157impl Printable for StructType {
158    fn fmt(
159        &self,
160        ctx: &Context,
161        state: &printable::State,
162        f: &mut core::fmt::Formatter<'_>,
163    ) -> core::fmt::Result {
164        write!(f, "<")?;
165
166        use core::cell::RefCell;
167        // Ugly, but also the simplest way to avoid infinite recursion.
168        // MLIR does the same: see LLVMTypeSyntax::printStructType.
169        thread_local! {
170            // We use a vec instead of a HashMap hoping that this isn't
171            // going to be large, in which case vec would be faster.
172            static IN_PRINTING: RefCell<Vec<Identifier>>  = const { RefCell::new(vec![]) };
173        }
174        if let Some(name) = &self.name {
175            let in_printing = IN_PRINTING.with(|f| f.borrow().contains(name));
176            if in_printing {
177                return write!(f, "{}>", name.clone());
178            }
179            IN_PRINTING.with(|f| f.borrow_mut().push(name.clone()));
180            write!(f, "{name}")?;
181            if !self.is_opaque() {
182                write!(f, " ")?;
183            }
184        }
185
186        if let Some(fields) = &self.fields {
187            enclosed(
188                "{ ",
189                " }",
190                list_with_sep(fields, ListSeparator::CharSpace(',')),
191            )
192            .fmt(ctx, state, f)?;
193        }
194
195        // Done processing this struct. Remove it from the stack.
196        if let Some(name) = &self.name {
197            assert!(IN_PRINTING.with(|f| f.borrow().last().unwrap() == name));
198            IN_PRINTING.with(|f| f.borrow_mut().pop());
199        }
200        write!(f, ">")
201    }
202}
203
204impl Hash for StructType {
205    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
206        match &self.name {
207            Some(name) => name.hash(state),
208            None => self
209                .fields
210                .as_ref()
211                .expect("Anonymous struct must have its fields set")
212                .iter()
213                .for_each(|field_type| {
214                    field_type.hash(state);
215                }),
216        }
217    }
218}
219
220impl PartialEq for StructType {
221    fn eq(&self, other: &Self) -> bool {
222        match (&self.name, &other.name) {
223            (Some(name), Some(other_name)) => name == other_name,
224            (None, None) => self.fields == other.fields,
225            _ => false,
226        }
227    }
228}
229
230impl Parsable for StructType {
231    type Arg = ();
232    type Parsed = TypedHandle<Self>;
233
234    fn parse<'a>(
235        state_stream: &mut StateStream<'a>,
236        _arg: Self::Arg,
237    ) -> ParseResult<'a, Self::Parsed>
238    where
239        Self: Sized,
240    {
241        let body_parser = || {
242            // Parse multiple type annotated fields separated by ',', all of it delimited by braces.
243            delimited_list_parser('{', '}', ',', type_parser())
244        };
245
246        let named = spaced((location(), Identifier::parser(())))
247            .and(spaced(optional(body_parser())))
248            .map(|((loc, name), body_opt)| (loc, Some(name), body_opt));
249        let anonymous = spaced((location(), body_parser()))
250            .map(|(loc, body)| (loc, None::<Identifier>, Some(body)));
251
252        // A struct type is named or anonymous.
253        let mut struct_parser = between(token('<'), token('>'), named.or(anonymous));
254
255        let (loc, name_opt, body_opt) = struct_parser.parse_stream(state_stream).into_result()?.0;
256        let ctx = &mut state_stream.state.ctx;
257        if let Some(name) = name_opt {
258            StructType::get_named(ctx, name, body_opt)
259                .map_err(|mut err| {
260                    err.set_loc(loc);
261                    err
262                })
263                .into_parse_result()
264        } else {
265            Ok(StructType::get_unnamed(
266                ctx,
267                body_opt.expect("Without a name, a struct type must have a body."),
268            ))
269            .into_parse_result()
270        }
271    }
272}
273
274impl Eq for StructType {}
275
276/// A pointer, corresponding to LLVM's pointer type. It is opaque (carries no
277/// pointee type) but carries an address space. The address space is always
278/// printed, e.g. `llvm.ptr (0)` or `llvm.ptr (1)`.
279#[pliron_type(
280    name = "llvm.ptr",
281    generate_get = true,
282    format = "`(` $address_space `)`",
283    verifier = "succ"
284)]
285#[derive(Hash, PartialEq, Eq, Debug)]
286pub struct PointerType {
287    address_space: u32,
288}
289
290impl PointerType {
291    /// The address space of this pointer.
292    pub fn address_space(&self) -> u32 {
293        self.address_space
294    }
295}
296
297/// Array type, corresponding to LLVM's array type.
298#[pliron_type(
299    name = "llvm.array",
300    generate_get = true,
301    format = "`[` $size ` x ` $elem `]`",
302    verifier = "succ"
303)]
304#[derive(Hash, PartialEq, Eq, Debug)]
305pub struct ArrayType {
306    elem: TypeHandle,
307    size: u64,
308}
309
310impl ArrayType {
311    /// Get array element type.
312    pub fn elem_type(&self) -> TypeHandle {
313        self.elem
314    }
315
316    /// Get array size.
317    pub fn size(&self) -> u64 {
318        self.size
319    }
320}
321
322#[pliron_type(name = "llvm.void", generate_get = true, format, verifier = "succ")]
323#[derive(Hash, PartialEq, Eq, Debug)]
324pub struct VoidType;
325
326#[pliron_type(
327    name = "llvm.func",
328    generate_get = true,
329    format = "`<` $res `(` vec($args, CharSpace(`,`)) `) variadic = ` $is_var_arg `>`",
330    verifier = "succ"
331)]
332#[derive(Hash, PartialEq, Eq, Debug)]
333pub struct FuncType {
334    res: TypeHandle,
335    args: Vec<TypeHandle>,
336    is_var_arg: bool,
337}
338
339#[derive(Debug, Error)]
340pub enum FuncTypeErr {
341    #[error("Expected at most one result")]
342    TooManyResults,
343}
344
345impl FuncType {
346    /// Result type
347    pub fn result_type(&self) -> TypeHandle {
348        self.res
349    }
350
351    /// Is this a variadic function type?
352    pub fn is_var_arg(&self) -> bool {
353        self.is_var_arg
354    }
355}
356
357#[type_interface_impl]
358impl FunctionTypeInterface for FuncType {
359    fn arg_types(&self) -> Vec<TypeHandle> {
360        self.args.clone()
361    }
362    fn res_types(&self) -> Vec<TypeHandle> {
363        vec![self.res]
364    }
365}
366
367/// Kind of vector type: fixed or scalable.
368/// See LLVM language reference for semantic details.
369#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
370#[format]
371pub enum VectorTypeKind {
372    Fixed,
373    Scalable,
374}
375
376#[pliron_type(
377    name = "llvm.vector",
378    generate_get = true,
379    format = "`<` $kind ` x ` $num_elems ` x ` $elem_ty `>`",
380    verifier = "succ"
381)]
382#[derive(Hash, PartialEq, Eq, Debug)]
383pub struct VectorType {
384    elem_ty: TypeHandle,
385    num_elems: u32,
386    kind: VectorTypeKind,
387}
388
389impl VectorType {
390    /// Get the element type.
391    pub fn elem_type(&self) -> TypeHandle {
392        self.elem_ty
393    }
394
395    /// Get the number of elements.
396    pub fn num_elements(&self) -> u32 {
397        self.num_elems
398    }
399
400    /// Is this a scalable vector type?
401    pub fn is_scalable(&self) -> bool {
402        self.kind == VectorTypeKind::Scalable
403    }
404
405    /// Get the scalable/fixed kind of this vector type.
406    pub fn kind(&self) -> VectorTypeKind {
407        self.kind
408    }
409}
410
411#[cfg(test)]
412mod tests {
413
414    use crate::types::{FuncType, PointerType, StructType, VoidType};
415    use expect_test::expect;
416    use pliron::{
417        builtin::types::{IntegerType, Signedness},
418        combine::{self, Parser, eof, token},
419        context::Context,
420        derive::{pliron_type, verify_succ},
421        identifier::Identifier,
422        irfmt::parsers::{spaced, type_parser},
423        parsable::{Parsable, ParseResult, StateStream, parse_from_str},
424        printable::{self, Printable},
425        result::{ExpectOk, Result},
426        r#type::{TypeHandle, TypedHandle},
427    };
428
429    #[test]
430    fn test_struct() -> Result<()> {
431        let ctx = Context::new();
432        let int64 = IntegerType::get(&ctx, 64, Signedness::Signless).into();
433        let linked_list_id: Identifier = "LinkedList".try_into().unwrap();
434
435        // Create an opaque struct since we want a recursive type.
436        let list_struct: TypeHandle =
437            StructType::get_named(&ctx, linked_list_id.clone(), None)?.into();
438        assert!(
439            list_struct
440                .deref(&ctx)
441                .downcast_ref::<StructType>()
442                .unwrap()
443                .is_opaque()
444        );
445        let list_struct_ptr = TypedPointerType::get(&ctx, list_struct).into();
446        let fields = vec![int64, list_struct_ptr];
447        // Set the struct body now.
448        StructType::get_named(&ctx, linked_list_id.clone(), Some(fields))?;
449        assert!(
450            !list_struct
451                .deref(&ctx)
452                .downcast_ref::<StructType>()
453                .unwrap()
454                .is_opaque()
455        );
456
457        let list_struct_2 = StructType::get_named(&ctx, linked_list_id, None)
458            .unwrap()
459            .into();
460        assert!(list_struct == list_struct_2);
461
462        assert_eq!(
463            list_struct.disp(&ctx).to_string(),
464            "llvm.struct <LinkedList { builtin.integer i64, llvm.typed_ptr <llvm.struct <LinkedList>> }>"
465        );
466
467        let head_fields = vec![int64, list_struct_ptr];
468        let head_struct = StructType::get_unnamed(&ctx, head_fields.clone());
469        let head_struct2 = StructType::get_unnamed(&ctx, head_fields);
470        assert!(head_struct == head_struct2);
471
472        Ok(())
473    }
474
475    /// A pointer type that knows the type it points to.
476    /// This used to be in LLVM earlier, but the latest version
477    /// is now type-erased (https://llvm.org/docs/OpaquePointers.html)
478    #[verify_succ]
479    #[pliron_type(name = "llvm.typed_ptr", generate_get = true)]
480    #[derive(Hash, PartialEq, Eq, Debug)]
481    pub struct TypedPointerType {
482        to: TypeHandle,
483    }
484
485    impl TypedPointerType {
486        /// Get the pointee type.
487        pub fn get_pointee_type(&self) -> TypeHandle {
488            self.to
489        }
490    }
491
492    impl Printable for TypedPointerType {
493        fn fmt(
494            &self,
495            ctx: &Context,
496            _state: &printable::State,
497            f: &mut core::fmt::Formatter<'_>,
498        ) -> core::fmt::Result {
499            write!(f, "<{}>", self.to.disp(ctx))
500        }
501    }
502
503    impl Parsable for TypedPointerType {
504        type Arg = ();
505        type Parsed = TypedHandle<Self>;
506
507        fn parse<'a>(
508            state_stream: &mut StateStream<'a>,
509            _arg: Self::Arg,
510        ) -> ParseResult<'a, Self::Parsed>
511        where
512            Self: Sized,
513        {
514            combine::between(token('<'), token('>'), spaced(type_parser()))
515                .parse_stream(state_stream)
516                .map(|pointee_ty| TypedPointerType::get(state_stream.state.ctx, pointee_ty))
517                .into()
518        }
519    }
520
521    #[test]
522    fn test_pointer_types() {
523        let ctx = Context::new();
524        let int32_1 = IntegerType::get(&ctx, 32, Signedness::Signed);
525        let int64 = IntegerType::get(&ctx, 64, Signedness::Signed).into();
526
527        let int64pointer = TypedPointerType::get(&ctx, int64);
528        assert_eq!(
529            int64pointer.disp(&ctx).to_string(),
530            "llvm.typed_ptr <builtin.integer si64>"
531        );
532        assert!(int64pointer == TypedPointerType::get(&ctx, int64));
533
534        assert!(
535            int64
536                .deref(&ctx)
537                .downcast_ref::<IntegerType>()
538                .unwrap()
539                .width()
540                == 64
541        );
542
543        assert!(IntegerType::get(&ctx, 32, Signedness::Signed) == int32_1);
544        assert!(TypedPointerType::get(&ctx, int64) == int64pointer);
545        assert!(int64pointer.deref(&ctx).get_pointee_type() == int64);
546    }
547
548    #[test]
549    fn test_pointer_type_parsing() {
550        let mut ctx = Context::new();
551
552        let res = parse_from_str(
553            type_parser(),
554            &mut ctx,
555            "llvm.typed_ptr <builtin.integer si64>",
556        )
557        .expect_ok(&ctx);
558        assert_eq!(
559            &res.disp(&ctx).to_string(),
560            "llvm.typed_ptr <builtin.integer si64>"
561        );
562    }
563
564    #[test]
565    fn test_opaque_pointer_addrspace() {
566        let mut ctx = Context::new();
567
568        // The address space is always printed, so addrspace 0 round-trips as
569        // `llvm.ptr (0)`.
570        let res = parse_from_str(type_parser(), &mut ctx, "llvm.ptr (0)").expect_ok(&ctx);
571        assert_eq!(res.disp(&ctx).to_string().trim(), "llvm.ptr (0)");
572        assert_eq!(
573            res.deref(&ctx)
574                .downcast_ref::<PointerType>()
575                .unwrap()
576                .address_space(),
577            0
578        );
579
580        // A non-zero address space round-trips as `llvm.ptr (N)`.
581        let res = parse_from_str(type_parser(), &mut ctx, "llvm.ptr (3)").expect_ok(&ctx);
582        assert_eq!(res.disp(&ctx).to_string().trim(), "llvm.ptr (3)");
583        assert_eq!(
584            res.deref(&ctx)
585                .downcast_ref::<PointerType>()
586                .unwrap()
587                .address_space(),
588            3
589        );
590    }
591
592    #[test]
593    fn test_fp16_type_roundtrip() {
594        let mut ctx = Context::new();
595        let res = parse_from_str(type_parser(), &mut ctx, "builtin.fp16").expect_ok(&ctx);
596        assert_eq!(res.disp(&ctx).to_string().trim(), "builtin.fp16");
597        assert!(
598            res.deref(&ctx)
599                .downcast_ref::<pliron::builtin::types::FP16Type>()
600                .is_some()
601        );
602    }
603
604    #[test]
605    fn test_struct_type_parsing() {
606        let mut ctx = Context::new();
607
608        let res = parse_from_str(
609            type_parser(),
610            &mut ctx,
611            "llvm.struct <LinkedList { builtin.integer i64, llvm.typed_ptr <llvm.struct <LinkedList>> }>",
612        )
613        .expect_ok(&ctx);
614        assert_eq!(
615            &res.disp(&ctx).to_string(),
616            "llvm.struct <LinkedList { builtin.integer i64, llvm.typed_ptr <llvm.struct <LinkedList>> }>"
617        );
618
619        // Test parsing an opaque struct.
620        let test_string = "llvm.struct <ExternStruct>";
621        let res = parse_from_str(type_parser(), &mut ctx, test_string).expect_ok(&ctx);
622        assert_eq!(&res.disp(&ctx).to_string(), test_string);
623        {
624            let res = res.deref(&ctx);
625            let res = res.downcast_ref::<StructType>().unwrap();
626            assert!(res.is_opaque() && res.is_named());
627        }
628
629        // Test parsing an unnamed struct.
630        let test_string = "llvm.struct <{ builtin.integer i8 }>";
631        let res = parse_from_str(type_parser(), &mut ctx, test_string).expect_ok(&ctx);
632        assert_eq!(&res.disp(&ctx).to_string(), test_string);
633        {
634            let res = res.deref(&ctx);
635            let res = res.downcast_ref::<StructType>().unwrap();
636            assert!(!res.is_opaque() && !res.is_named());
637        }
638    }
639
640    #[test]
641    fn test_struct_type_errs() {
642        let mut ctx = Context::new();
643
644        let _ = parse_from_str(
645            type_parser(),
646            &mut ctx,
647            "llvm.struct < My1 { builtin.integer i8 } >",
648        )
649        .expect_ok(&ctx);
650
651        let err_msg = format!(
652            "{}",
653            parse_from_str(
654                type_parser(),
655                &mut ctx,
656                "llvm.struct < My1 { builtin.integer i16 } >",
657            )
658            .unwrap_err()
659        );
660
661        let expected_err_msg = expect![[r#"
662            Compilation error: invalid input program.
663            Parse error at line: 1, column: 15
664            struct My1 already exists and is different
665        "#]];
666        expected_err_msg.assert_eq(&err_msg);
667    }
668
669    #[test]
670    fn test_functype_parsing() {
671        let mut ctx = Context::new();
672
673        let si32 = IntegerType::get(&ctx, 32, Signedness::Signed);
674
675        let input = "llvm.func <llvm.void (builtin.integer si32) variadic = false>";
676        let res = parse_from_str(type_parser().and(eof()), &mut ctx, input)
677            .expect_ok(&ctx)
678            .0;
679
680        let void_ty = VoidType::get(&ctx);
681        assert!(res == FuncType::get(&ctx, void_ty.to_handle(), vec![si32.into()], false).into());
682        assert_eq!(input, &res.disp(&ctx).to_string());
683    }
684}
pliron_llvm/types.rs

pliron_llvm/
types.rs
