Implement type inference for rest attributes

crates/ide/src/ty/display.rs

@@ -75,15 +75,18 @@ impl fmt::Display for TyDisplay<'_> {
                 Ok(())
             }
             Ty::Attrset(set) => {
-                if config.max_attrset_depth == 0 {
+                if config.max_attrset_depth == 0 || config.max_attrset_fields == 0 {
                     return "{…}".fmt(f);
                 }
                 config.max_attrset_depth -= 1;
                 config.lambda_need_parentheses = false;
 
                 "{".fmt(f)?;
                 let mut first = true;
-                for (name, ty, _src) in set.iter().take(config.max_attrset_fields) {
+                let max_fields = config
+                    .max_attrset_fields
+                    .saturating_sub(set.rest.is_some() as usize);
+                for (name, ty, _src) in set.iter().take(max_fields) {
                     if first {
                         first = false;
                     } else {
@@ -93,7 +96,16 @@ impl fmt::Display for TyDisplay<'_> {
                     let value = Self { ty, config };
                     write!(f, " {name}: {value}")?;
                 }
-                if set.len() > config.max_attrset_fields {
+                if let Some(rest) = &set.rest {
+                    if !first {
+                        ",".fmt(f)?;
+                    }
+                    let value = Self {
+                        ty: &rest.0,
+                        config,
+                    };
+                    write!(f, " …: {value} }}")
+                } else if set.len() > max_fields {
                     ", … }".fmt(f)
                 } else {
                     " }".fmt(f)
@@ -102,3 +114,43 @@ impl fmt::Display for TyDisplay<'_> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use expect_test::{expect, Expect};
+
+    use super::{Config, TyDisplay};
+    use crate::ty::Ty;
+
+    #[track_caller]
+    fn check_max_fields(max_attrset_fields: usize, ty: &Ty, expect: Expect) {
+        let disp = TyDisplay {
+            ty,
+            config: Config {
+                max_attrset_fields,
+                ..super::Config::FULL
+            },
+        };
+        expect.assert_eq(&disp.to_string());
+    }
+
+    #[test]
+    fn attrset() {
+        let ty = &ty!({ "a": int, "b": string, "c": bool });
+        check_max_fields(0, ty, expect!["{…}"]);
+        check_max_fields(1, ty, expect!["{ a: int, … }"]);
+        check_max_fields(2, ty, expect!["{ a: int, b: string, … }"]);
+        check_max_fields(3, ty, expect!["{ a: int, b: string, c: bool }"]);
+        check_max_fields(4, ty, expect!["{ a: int, b: string, c: bool }"]);
+    }
+
+    #[test]
+    fn attrset_rest() {
+        let ty = &ty!({ "a": int, "b": string, _: bool });
+        check_max_fields(0, ty, expect!["{…}"]);
+        check_max_fields(1, ty, expect!["{ …: bool }"]);
+        check_max_fields(2, ty, expect!["{ a: int, …: bool }"]);
+        check_max_fields(3, ty, expect!["{ a: int, b: string, …: bool }"]);
+        check_max_fields(4, ty, expect!["{ a: int, b: string, …: bool }"]);
+    }
+}

crates/ide/src/ty/infer.rs

@@ -37,6 +37,7 @@ enum Ty {
 
     List(TyVar),
     Lambda(TyVar, TyVar),
+    // TODO: Add support for `rest` similar to super::Attrset.
     Attrset(Attrset),
 
     External(super::Ty),
@@ -532,13 +533,13 @@ impl<'a> Collector<'a> {
                 let b = self.collect(b);
                 super::Ty::Lambda(a.into(), b.into())
             }
-            Ty::Attrset(set) => {
-                let set = set
+            Ty::Attrset(fields) => {
+                let fields = fields
                     .0
                     .into_iter()
                     .map(|(name, (ty, src))| (name, self.collect(ty), src))
                     .collect();
-                super::Ty::Attrset(super::Attrset(set))
+                super::Ty::Attrset(super::Attrset { fields, rest: None })
             }
             Ty::External(ty) => ty,
         }

crates/ide/src/ty/known.rs

@@ -57,15 +57,18 @@ fn merge_attrset(lhs: &Ty, rhs: &Ty) -> Ty {
     let rhs = rhs.as_attrset().unwrap();
     // Put the RHS on the front and ...
     let mut xs = rhs
-        .0
+        .fields
         .iter()
-        .chain(lhs.0.iter())
+        .chain(lhs.fields.iter())
         .map(|(name, ty, src)| (name.clone(), ty.clone(), *src))
         .collect::<Vec<_>>();
     // ... run stable sort to prefer RHS when duplicated.
     xs.sort_by(|(lhs, ..), (rhs, ..)| lhs.cmp(rhs));
     xs.dedup_by(|(lhs, ..), (rhs, ..)| lhs == rhs);
-    Ty::Attrset(Attrset(xs.into()))
+    Ty::Attrset(Attrset {
+        fields: xs.into(),
+        rest: rhs.rest.clone(),
+    })
 }
 
 /// https://nixos.wiki/wiki/Flakes
@@ -87,6 +90,7 @@ pub fn flake(inputs: &[&str]) -> Ty {
             .iter()
             .copied()
             .map(|name| (name, input_ty.clone(), AttrSource::Unknown)),
+        None,
     ));
 
     let outputs_param_ty = Ty::Attrset(Attrset::from_internal(
@@ -95,6 +99,7 @@ pub fn flake(inputs: &[&str]) -> Ty {
             .copied()
             .chain(Some("self"))
             .map(|name| (name, FLAKE.clone(), AttrSource::Unknown)),
+        None,
     ));
 
     ty!({

crates/ide/src/ty/mod.rs

@@ -28,18 +28,37 @@ macro_rules! ty {
 
     (($($inner:tt)*)) => { ty!($($inner)*) };
     ([$($inner:tt)*]) => { $crate::ty::Ty::List(::std::sync::Arc::new(ty!($($inner)*)))};
-    ({ $($key:literal : $ty:tt),* $(,)? $(_ : $rest_ty:tt)? }) => {{
-        // TODO: Rest type.
-        $(let _ = ty!($rest_ty);)?
-        $crate::ty::Ty::Attrset($crate::ty::Attrset::from_internal([
-            $(($key, ty!($ty), $crate::ty::AttrSource::Unknown),)*
-        ]))
+    ({ $($key:literal : $ty:tt),* $(,)? }) => {{
+        $crate::ty::Ty::Attrset($crate::ty::Attrset::from_internal(
+            [
+                $(($key, ty!($ty), $crate::ty::AttrSource::Unknown),)*
+            ],
+            None,
+        ))
     }};
-    ({($src:expr) $($key:literal : $ty:tt),* $(,)? $(_ : $rest_ty:tt)? }) => {{
-        $(let _ = ty!($rest_ty);)?
-        $crate::ty::Ty::Attrset($crate::ty::Attrset::from_internal([
-            $(($key, ty!($ty), $src),)*
-        ]))
+    ({ $($key:literal : $ty:tt),* $(,)? _ : $rest_ty:tt }) => {{
+        $crate::ty::Ty::Attrset($crate::ty::Attrset::from_internal(
+            [
+                $(($key, ty!($ty), $crate::ty::AttrSource::Unknown),)*
+            ],
+            Some((ty!($rest_ty), $crate::ty::AttrSource::Unknown)),
+        ))
+    }};
+    ({($src:expr) $($key:literal : $ty:tt),* $(,)? }) => {{
+        $crate::ty::Ty::Attrset($crate::ty::Attrset::from_internal(
+            [
+                $(($key, ty!($ty), $src),)*
+            ],
+            None
+        ))
+    }};
+    ({($src:expr) $($key:literal : $ty:tt),* $(,)? _ : $rest_ty:tt }) => {{
+        $crate::ty::Ty::Attrset($crate::ty::Attrset::from_internal(
+            [
+                $(($key, ty!($ty), $src),)*
+            ],
+            Some((ty!($rest_ty), $src))
+        ))
     }};
     ($arg:tt -> $($ret:tt)*) => {
         $crate::ty::Ty::Lambda(
@@ -122,11 +141,17 @@ impl fmt::Debug for Ty {
 
 // Invariant: sorted by names.
 #[derive(Debug, Clone, PartialEq, Eq)]
-pub struct Attrset(Arc<[(SmolStr, Ty, AttrSource)]>);
+pub struct Attrset {
+    fields: Arc<[(SmolStr, Ty, AttrSource)]>,
+    rest: Option<Arc<(Ty, AttrSource)>>,
+}
 
 impl Default for Attrset {
     fn default() -> Self {
-        Self(Arc::new([]))
+        Self {
+            fields: Arc::new([]),
+            rest: None,
+        }
     }
 }
 
@@ -136,32 +161,41 @@ impl Attrset {
     /// # Panics
     /// The given iterator must have no duplicated fields, or it'll panic.
     #[track_caller]
-    pub fn from_internal<'a>(iter: impl IntoIterator<Item = (&'a str, Ty, AttrSource)>) -> Self {
-        let mut set = iter
+    pub fn from_internal<'a>(
+        iter: impl IntoIterator<Item = (&'a str, Ty, AttrSource)>,
+        rest: Option<(Ty, AttrSource)>,
+    ) -> Self {
+        let mut fields = iter
             .into_iter()
             .map(|(name, ty, src)| (SmolStr::from(name), ty, src))
             .collect::<Arc<[_]>>();
-        Arc::get_mut(&mut set)
+        Arc::get_mut(&mut fields)
             .unwrap()
             .sort_by(|(lhs, ..), (rhs, ..)| lhs.cmp(rhs));
         assert!(
-            set.windows(2).all(|w| w[0].0 != w[1].0),
+            fields.windows(2).all(|w| w[0].0 != w[1].0),
             "Duplicated fields",
         );
-        Self(set)
+        Self {
+            fields,
+            rest: rest.map(Arc::new),
+        }
     }
 
     pub fn is_empty(&self) -> bool {
-        self.0.is_empty()
+        self.fields.is_empty()
     }
 
     pub fn len(&self) -> usize {
-        self.0.len()
+        self.fields.len()
     }
 
     fn get_all(&self, field: &str) -> Option<(&Ty, AttrSource)> {
-        let i = self.0.binary_search_by(|p| (*p.0).cmp(field)).ok()?;
-        Some((&self.0[i].1, self.0[i].2))
+        if let Ok(i) = self.fields.binary_search_by(|p| (*p.0).cmp(field)) {
+            Some((&self.fields[i].1, self.fields[i].2))
+        } else {
+            self.rest.as_ref().map(|rest| (&rest.0, rest.1))
+        }
     }
 
     pub fn get(&self, field: &str) -> Option<&Ty> {
@@ -173,7 +207,7 @@ impl Attrset {
     }
 
     pub fn iter(&self) -> impl Iterator<Item = (&SmolStr, &Ty, AttrSource)> + '_ {
-        self.0.iter().map(|(k, ty, src)| (k, ty, *src))
+        self.fields.iter().map(|(k, ty, src)| (k, ty, *src))
     }
 }
 

crates/ide/src/ty/tests.rs

@@ -228,8 +228,28 @@ fn flake_file() {
     );
 }
 
+#[test]
+fn rest_type() {
+    check_name(
+        "bar",
+        r"
+#- /flake.nix
+{
+    outputs = { self }: {
+        apps.x86_64-linux = {
+            foo = let bar = bar; in bar;
+        };
+    };
+}
+        ",
+        expect!["{ program: string, type: string }"],
+    );
+}
+
 #[test]
 fn builtins() {
     check("true", expect!["bool"]);
     check("builtins.length [ ]", expect!["int"]);
+    check("builtins.readDir ./.", expect!["{ …: string }"]);
+    check("(builtins.readDir ./.).foo", expect!["string"]);
 }