Add support for struct and field projection

tiif · tiif · commit bc5cb39fc141 · 2025-08-04T10:26:53.000+02:00
diff --git a/crates/formality-check/src/mini_rust_check.rs b/crates/formality-check/src/mini_rust_check.rs
@@ -3,15 +3,15 @@ use std::iter::zip;
 use formality_core::{Fallible, Map, Upcast};
 use formality_prove::Env;
 use formality_rust::grammar::minirust::ArgumentExpression::{ByValue, InPlace};
-use formality_rust::grammar::minirust::PlaceExpression::Local;
-use formality_rust::grammar::minirust::ValueExpression::{Constant, Fn, Load};
+use formality_rust::grammar::minirust::PlaceExpression::*;
+use formality_rust::grammar::minirust::ValueExpression::{Constant, Fn, Load, Struct};
 use formality_rust::grammar::minirust::{
     self, ty_is_int, ArgumentExpression, BasicBlock, BbId, LocalId, PlaceExpression,
     ValueExpression,
 };
-use formality_rust::grammar::FnBoundData;
-use formality_types::grammar::{CrateId, FnId};
-use formality_types::grammar::{Relation, Ty, Wcs};
+use formality_rust::grammar::{FnBoundData, StructBoundData};
+use formality_types::grammar::{AdtId, Relation, Ty, TyData, Wcs};
+use formality_types::grammar::{CrateId, FnId, RigidName};
 
 use crate::{Check, CrateItem};
 use anyhow::bail;
@@ -95,6 +95,7 @@ impl Check<'_> {
             callee_input_tys: Map::new(),
             crate_id: crate_id.clone(),
             fn_args: body.args.clone(),
+            adt_tys: Map::new(),
         };
 
         // (4) Check statements in body are valid
@@ -276,6 +277,38 @@ impl Check<'_> {
                 };
                 place_ty = ty;
             }
+            Field(field_projection) => {
+                let Local(ref local_id) = *field_projection.root else {
+                    bail!("Only Local is allowed as the root of FieldProjection")
+                };
+
+                // Check if the index is valid for the tuple.
+                // FIXME: use let chain here?
+
+                let Some(ty) = env.local_variables.get(&local_id) else {
+                    bail!("The local id used in PlaceExpression::Field is invalid.")
+                };
+
+                // Get the ADT type information.
+                let TyData::RigidTy(rigid_ty) = ty.data() else {
+                    bail!("The type for field projection must be rigid ty")
+                };
+
+                // FIXME: it'd be nice to have the field information in ty data
+                let RigidName::AdtId(ref adt_id) = rigid_ty.name else {
+                    bail!("The type for field projection must be adt")
+                };
+
+                let Some(tys) = env.adt_tys.get(&adt_id) else {
+                    bail!("The ADT used is invalid.")
+                };
+
+                if field_projection.index >= tys.len() {
+                    bail!("The field index used in PlaceExpression::Field is invalid.")
+                }
+
+                place_ty = tys[field_projection.index].clone();
+            }
         }
         Ok(place_ty.clone())
     }
@@ -317,6 +350,7 @@ impl Check<'_> {
                             if fn_declared.id == *fn_id {
                                 let fn_bound_data =
                                     typeck_env.env.instantiate_universally(&fn_declared.binder);
+                                // FIXME: maybe we should store the information somewhere else, like in the value expression?
                                 // Store the callee information in typeck_env, we will need this when type checking Terminator::Call.
                                 typeck_env
                                     .callee_input_tys
@@ -341,6 +375,77 @@ impl Check<'_> {
                 // it will be rejected by the parser.
                 Ok(constant.get_ty())
             }
+            Struct(value_expressions, ty) => {
+                let mut struct_field_ty = Vec::new();
+
+                // Check if the adt type is valid in current crate.
+                if let TyData::RigidTy(rigid_ty) = ty.data() {
+                    if let RigidName::AdtId(adt_id) = &rigid_ty.name {
+                        // Find the crate that is currently being typeck.
+                        let curr_crate = self
+                            .program
+                            .crates
+                            .iter()
+                            .find(|c| c.id == typeck_env.crate_id)
+                            .unwrap();
+
+                        // Find the struct from current crate.
+                        let target_struct = curr_crate.items.iter().find(|item| {
+                            match item {
+                                CrateItem::Struct(struct_item) => {
+                                    if struct_item.id == *adt_id {
+                                        // Get the ty data of the field.
+                                        let (
+                                            _,
+                                            StructBoundData {
+                                                where_clauses: _,
+                                                fields,
+                                            },
+                                        ) = struct_item.binder.open();
+                                        for field in fields {
+                                            struct_field_ty.push(field.ty);
+                                        }
+                                        return true;
+                                    }
+                                    false
+                                }
+                                _ => false,
+                            }
+                        });
+
+                        if target_struct.is_none() {
+                            bail!("The type used in Tuple is not declared in current crate")
+                        }
+                        // We will need the adt type information when type checking field projection.
+                        typeck_env
+                            .adt_tys
+                            .insert(adt_id.clone(), struct_field_ty.clone());
+                    }
+                }
+
+                // Make sure the length of value expression matches the length of field of adt.
+                if value_expressions.len() != struct_field_ty.len() {
+                    bail!("The length of ValueExpression::Tuple does not match the type of the ADT declared")
+                }
+
+                let expression_ty_pair = zip(value_expressions, struct_field_ty);
+
+                // FIXME: we only support const in value expression of tuple for now, we can add support
+                // more in future.
+
+                for (value_expression, declared_ty) in expression_ty_pair {
+                    let Constant(_) = value_expression else {
+                        bail!("Only Constant is supported in ValueExpression::Tuple for now.")
+                    };
+                    let ty = self.check_value(typeck_env, value_expression)?;
+
+                    // Make sure the type matches the declared adt.
+                    if ty != declared_ty {
+                        bail!("The type in ValueExpression::Tuple does not match the ADT declared")
+                    }
+                }
+                Ok(ty.clone())
+            }
         }
     }
 
@@ -402,4 +507,7 @@ struct TypeckEnv {
 
     /// LocalId of function argument.
     fn_args: Vec<LocalId>,
+
+    /// Type information of adt
+    adt_tys: Map<AdtId, Vec<Ty>>,
 }
diff --git a/crates/formality-rust/src/grammar/minirust.rs b/crates/formality-rust/src/grammar/minirust.rs
@@ -5,10 +5,13 @@ use formality_types::grammar::{Parameter, RigidName, ScalarId, Ty, TyData};
 use crate::grammar::minirust::ConstTypePair::*;
 use crate::grammar::FnId;
 
+use std::sync::Arc;
+
 // This definition is based on [MiniRust](https://github.com/minirust/minirust/blob/master/spec/lang/syntax.md).
 
 id!(BbId);
 id!(LocalId);
+id!(FieldId);
 
 // Example:
 //
@@ -138,8 +141,11 @@ pub enum ValueExpression {
     Constant(ConstTypePair),
     #[grammar(fn_id $v0)]
     Fn(FnId),
-    // #[grammar($(v0) as $v1)]
-    // Tuple(Vec<ValueExpression>, Ty),
+    // FIXME: minirust uses typle to represent  arrays, structs, tuples (including unit).
+    // But I think it will be quite annoying to do typecking when we have all these types
+    // together, so I added a variant just for struct.
+    #[grammar(struct ${v0} as $v1)] // FIXME: use comma separated
+    Struct(Vec<ValueExpression>, Ty),
     // Union
     // Variant
     // GetDiscriminant
@@ -229,7 +235,17 @@ pub enum PlaceExpression {
     #[grammar(local($v0))]
     Local(LocalId),
     // Deref(Arc<ValueExpression>),
-    // Field(Arc<PlaceExpression>, FieldId),
+    // Project to a field.
+    #[grammar($v0)]
+    Field(FieldProjection),
     // Index
     // Downcast
 }
+
+#[term($root.$index)]
+pub struct FieldProjection {
+    /// The place to base the projection on.
+    pub root: Arc<PlaceExpression>,
+    /// The field to project to.
+    pub index: usize,
+}
diff --git a/src/test/mir_fn_bodies.rs b/src/test/mir_fn_bodies.rs
@@ -254,6 +254,37 @@ fn test_storage_live_dead() {
     )
 }
 
+/// Test valid program that uses struct.  
+#[test]
+fn test_struct() {
+    crate::assert_ok!(
+        [
+            crate Foo {
+                struct Dummy {
+                    value: u32,
+                }
+
+                fn foo (u32) -> u32 = minirust(v1) -> v0 {
+                    let v0: u32;
+                    let v1: u32;
+                    let v2: Dummy;
+
+                    bb0: {
+                        statements {
+                            local(v0) = load(local(v1));
+                            local(v2) = struct { constant(1: u32) } as Dummy;
+                            local(v2).0 = constant(2: u32);
+                        }
+                        return;
+                    }
+
+                };
+            }
+        ]
+        expect_test::expect![["()"]]
+    )
+}
+
 // Test what will happen if the next block does not exist for Terminator::Call.
 #[test]
 fn test_no_next_bb_for_call_terminator() {
@@ -737,3 +768,98 @@ fn test_fn_arg_storage_dead() {
         expect_test::expect!["Statement::StorageDead: trying to mark function arguments or return local as dead"]
     )
 }
+
+/// Test the behaviour of using invalid index for the struct field.
+#[test]
+fn test_invalid_struct_field() {
+    crate::assert_err!(
+        [
+            crate Foo {
+                struct Dummy {
+                    value: u32,
+                }
+
+                fn foo (u32) -> u32 = minirust(v1) -> v0 {
+                    let v0: u32;
+                    let v1: u32;
+                    let v2: Dummy;
+
+                    bb0: {
+                        statements {
+                            local(v0) = load(local(v1));
+                            local(v2) = struct { constant(1: u32) } as Dummy;
+                            local(v2).1 = constant(2: u32);
+                        }
+                        return;
+                    }
+
+                };
+            }
+        ]
+        []
+        expect_test::expect!["The field index used in PlaceExpression::Field is invalid."]
+    )
+}
+
+/// Test the behaviour of using non-adt local for field projection.
+#[test]
+fn test_field_projection_root_non_adt() {
+    crate::assert_err!(
+        [
+            crate Foo {
+                struct Dummy {
+                    value: u32,
+                }
+
+                fn foo (u32) -> u32 = minirust(v1) -> v0 {
+                    let v0: u32;
+                    let v1: u32;
+                    let v2: Dummy;
+
+                    bb0: {
+                        statements {
+                            local(v0) = load(local(v1));
+                            local(v2) = struct { constant(1: u32) } as Dummy;
+                            local(v1).1 = constant(2: u32);
+                        }
+                        return;
+                    }
+
+                };
+            }
+        ]
+        []
+        expect_test::expect!["The type for field projection must be adt"]
+    )
+}
+
+/// Test the behaviour of initialising the struct with wrong type.
+#[test]
+fn test_struct_wrong_type_in_initialisation() {
+    crate::assert_err!(
+        [
+            crate Foo {
+                struct Dummy {
+                    value: u32,
+                }
+
+                fn foo (u32) -> u32 = minirust(v1) -> v0 {
+                    let v0: u32;
+                    let v1: u32;
+                    let v2: Dummy;
+
+                    bb0: {
+                        statements {
+                            local(v0) = load(local(v1));
+                            local(v2) = struct { constant(false) } as Dummy;
+                        }
+                        return;
+                    }
+
+                };
+            }
+        ]
+        []
+        expect_test::expect!["The type in ValueExpression::Tuple does not match the ADT declared"]
+    )
+}
