Add input flake types to database inputs

crates/ide/src/base.rs

@@ -1,3 +1,4 @@
+use nix_interop::flake_output::FlakeOutput;
 use nix_interop::nixos_options::NixosOptions;
 use salsa::Durability;
 use std::collections::HashMap;
@@ -172,10 +173,12 @@ pub struct FlakeGraph {
     pub nodes: HashMap<SourceRootId, FlakeInfo>,
 }
 
+// FIXME: Make this a tree structure.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct FlakeInfo {
     pub flake_file: FileId,
     pub input_store_paths: HashMap<String, VfsPath>,
+    pub input_flake_outputs: HashMap<String, FlakeOutput>,
 }
 
 #[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]

crates/ide/src/def/tests.rs

@@ -169,6 +169,7 @@ fn source_root_flake() {
                 "nixpkgs".into(),
                 VfsPath::new("/nix/store/eeee"),
             )]),
+            input_flake_outputs: HashMap::new(),
         },
     );
 }

crates/ide/src/tests.rs

@@ -117,6 +117,7 @@ impl Fixture {
                     let flake_info = this.flake_info.insert(FlakeInfo {
                         flake_file: cur_file,
                         input_store_paths: HashMap::new(),
+                        input_flake_outputs: HashMap::new(),
                     });
                     for prop in iter {
                         if let Some((name, target)) = prop

crates/ide/src/ty/convert.rs

@@ -1,7 +1,11 @@
 //! Convert structures from Nix evaluation result into `Ty`s.
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use nix_interop::flake_output::{FlakeOutput, Type as OutputTy};
 use nix_interop::nixos_options::Ty as OptionTy;
 
-use crate::TyDatabase;
+use crate::{SourceRootId, TyDatabase};
 
 use super::{AttrSource, Attrset, Ty};
 
@@ -35,3 +39,29 @@ fn from_raw_ty(ty: &OptionTy) -> Ty {
         }
     }
 }
+
+pub(crate) fn flake_input_tys(db: &dyn TyDatabase, sid: SourceRootId) -> Arc<HashMap<String, Ty>> {
+    let Some(info) = db.source_root_flake_info(sid) else { return Arc::default() };
+    let tys = info
+        .input_flake_outputs
+        .iter()
+        .map(|(name, output)| (name.clone(), from_flake_output(output)))
+        .collect();
+    Arc::new(tys)
+}
+
+fn from_flake_output(out: &FlakeOutput) -> Ty {
+    match out {
+        FlakeOutput::Leaf(leaf) => match leaf.type_ {
+            OutputTy::NixosModule => ty!({}),
+            OutputTy::Derivation => ty!(derivation),
+            OutputTy::Unknown => ty!(?),
+        },
+        FlakeOutput::Attrset(set) => {
+            let fields = set
+                .iter()
+                .map(|(key, output)| (&**key, from_flake_output(output), AttrSource::Unknown));
+            Ty::Attrset(Attrset::from_internal(fields, None))
+        }
+    }
+}

crates/ide/src/ty/known.rs

@@ -41,7 +41,7 @@ pub static FETCH_TREE_RET: Lazy<Ty> = Lazy::new(|| {
     })
 });
 
-pub static FLAKE: Lazy<Ty> = Lazy::new(|| {
+pub static GENERIC_FLAKE: Lazy<Ty> = Lazy::new(|| {
     merge_attrset(
         &FETCH_TREE_RET,
         &ty!({
@@ -71,10 +71,9 @@ fn merge_attrset(lhs: &Ty, rhs: &Ty) -> Ty {
     })
 }
 
-/// <https://nixos.wiki/wiki/Flakes>
-pub fn flake(inputs: &[&str]) -> Ty {
-    // https://nixos.org/manual/nix/stable/command-ref/new-cli/nix3-flake.html#flake-references
-    let input_ty = merge_attrset(
+// https://nixos.org/manual/nix/stable/command-ref/new-cli/nix3-flake.html#flake-references
+static GENERIC_INPUT_DECL: Lazy<Ty> = Lazy::new(|| {
+    merge_attrset(
         &FETCH_TREE_ARG,
         &ty!({
             "inputs": {
@@ -83,31 +82,36 @@ pub fn flake(inputs: &[&str]) -> Ty {
                 }
             },
         }),
-    );
+    )
+});
 
-    let inputs_ty = Ty::Attrset(Attrset::from_internal(
-        inputs
-            .iter()
-            .copied()
-            .map(|name| (name, input_ty.clone(), AttrSource::Unknown)),
+/// <https://nixos.wiki/wiki/Flakes>
+pub fn flake(inputs: &[(&str, Ty)]) -> Ty {
+    let inputs_decl_ty = Ty::Attrset(Attrset::from_internal(
+        inputs.iter().map(|(name, ty)| {
+            let ty = merge_attrset(ty, &GENERIC_INPUT_DECL);
+            (*name, ty, AttrSource::Unknown)
+        }),
         None,
     ));
 
     let outputs_param_ty = Ty::Attrset(Attrset::from_internal(
         inputs
             .iter()
-            .copied()
             // Don't duplicate.
-            .filter(|name| *name != "self")
-            .chain(Some("self"))
-            .map(|name| (name, FLAKE.clone(), AttrSource::Unknown)),
+            .filter(|(name, _)| *name != "self")
+            .map(|(name, ty)| {
+                let ty = merge_attrset(ty, &GENERIC_FLAKE);
+                (*name, ty, AttrSource::Unknown)
+            })
+            .chain(Some(("self", ty!({}), AttrSource::Unknown))),
         None,
     ));
 
     ty!({
         "description": string,
         "nixConfig": { _: ? },
-        "inputs": (#inputs_ty),
+        "inputs": (#inputs_decl_ty),
         // https://nixos.org/manual/nix/unstable/command-ref/new-cli/nix3-develop.html?highlight=flake#flake-output-attributes
         "outputs": ((#outputs_param_ty) -> {
             "apps": {
@@ -285,7 +289,7 @@ fn builtins() -> Ty {
         "getAttr": (forall a, string -> { _: a } -> a),
         "getContext": (string -> { _: { "outputs": [string] } }),
         "getEnv": (string -> string),
-        "getFlake": (string -> (#FLAKE.clone())),
+        "getFlake": (string -> (#GENERIC_FLAKE.clone())),
         "groupBy": (forall a, (a -> string) -> [a] -> { _: [a] }),
         "hasAttr": (string -> { } -> bool),
         "hasContext": (string -> bool),

crates/ide/src/ty/mod.rs

@@ -78,7 +78,8 @@ mod union_find;
 mod tests;
 
 use crate::def::NameId;
-use crate::{DefDatabase, FileId, ModuleKind};
+use crate::{DefDatabase, FileId, ModuleKind, SourceRootId};
+use std::collections::HashMap;
 use std::fmt;
 use std::sync::Arc;
 
@@ -96,6 +97,9 @@ pub trait TyDatabase: DefDatabase {
 
     #[salsa::invoke(convert::options_to_config_ty)]
     fn nixos_config_ty(&self) -> Ty;
+
+    #[salsa::invoke(convert::flake_input_tys)]
+    fn flake_input_tys(&self, sid: SourceRootId) -> Arc<HashMap<String, Ty>>;
 }
 
 #[derive(Clone, PartialEq, Eq)]
@@ -234,13 +238,23 @@ fn module_expected_ty(db: &dyn TyDatabase, file: FileId) -> Option<Ty> {
             param_inputs,
             ..
         } => {
+            let sid = db.file_source_root(file);
+            let input_tys = db.flake_input_tys(sid);
             let mut inputs = explicit_inputs
                 .keys()
                 .chain(param_inputs.keys())
-                .map(|s| &**s)
+                .map(|s| {
+                    let input_ty = input_tys
+                        .get(&**s)
+                        .cloned()
+                        // NB. This must be an `Attrset`, so that `known::flake` will merge it
+                        // normally without panicking.
+                        .unwrap_or_else(|| Ty::Attrset(Attrset::default()));
+                    (&**s, input_ty)
+                })
                 .collect::<Vec<_>>();
-            inputs.sort();
-            inputs.dedup();
+            inputs.sort_by_key(|(name, _)| *name);
+            inputs.dedup_by_key(|(name, _)| *name);
             Some(known::flake(&inputs))
         }
         ModuleKind::Package { .. } => Some(known::PACKAGE.clone()),

crates/ide/src/ty/tests.rs

@@ -1,6 +1,12 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
 use crate::tests::TestDB;
-use crate::{DefDatabase, InferenceResult, Module, TyDatabase};
+use crate::{
+    DefDatabase, FlakeGraph, FlakeInfo, InferenceResult, Module, SourceDatabase, TyDatabase,
+};
 use expect_test::{expect, Expect};
+use nix_interop::flake_output::{FlakeOutput, Type};
 
 use super::Ty;
 
@@ -267,3 +273,48 @@ fn inputs_with_self() {
         expect!["int"],
     );
 }
+
+#[test]
+fn input_flake_ty() {
+    let src = r#"
+#- /flake.nix
+{
+    inputs.nixpkgs = "...";
+    outputs = { self, nixpkgs }: {
+        export = nixpkgs.hello;
+    };
+}
+    "#;
+
+    let nixpkgs_output = FlakeOutput::Attrset(HashMap::from_iter([(
+        "hello".into(),
+        FlakeOutput::Leaf(nix_interop::flake_output::Leaf {
+            type_: Type::Derivation,
+            name: None,
+            description: None,
+        }),
+    )]));
+
+    let expect = expect!["{ args: [string], builder: string, name: string, system: string }"];
+
+    let (mut db, file) = TestDB::single_file(src).unwrap();
+    let sid = db.file_source_root(file);
+    db.set_flake_graph(Arc::new(FlakeGraph {
+        nodes: HashMap::from_iter([(
+            sid,
+            FlakeInfo {
+                flake_file: file,
+                input_store_paths: HashMap::new(),
+                input_flake_outputs: HashMap::from_iter([("nixpkgs".into(), nixpkgs_output)]),
+            },
+        )]),
+    }));
+    let name = db
+        .module(file)
+        .names()
+        .find(|(_, n)| n.text == "export")
+        .expect("Name not found")
+        .0;
+    let ty = db.infer(file).ty_for_name(name).debug().to_string();
+    expect.assert_eq(&ty);
+}

crates/nil/src/server.rs

@@ -436,6 +436,7 @@ impl Server {
                 return Ok(Some(FlakeInfo {
                     flake_file,
                     input_store_paths: HashMap::new(),
+                    input_flake_outputs: HashMap::new(),
                 }));
             }
             Err(err) => {
@@ -456,6 +457,7 @@ impl Server {
         Ok(Some(FlakeInfo {
             flake_file,
             input_store_paths,
+            input_flake_outputs: HashMap::new(),
         }))
     }