Represent path segments in fhir

crates/flux-desugar/src/desugar.rs

@@ -1049,10 +1049,29 @@ trait DesugarCtxt<'genv, 'tcx: 'genv> {
 
     fn desugar_path(&mut self, path: &surface::Path) -> Result<fhir::Path<'genv>> {
         let res = self.resolver_output().path_res_map[&path.node_id];
-        let (args, bindings) = self.desugar_generic_args(res, &path.generics)?;
+        let [prefix @ .., last] = &path.segments[..] else {
+            span_bug!(path.span, "expected at least one segment")
+        };
+        let segments = self.genv().alloc_slice_with_capacity(
+            prefix.len() + 1,
+            prefix
+                .iter()
+                .map(|ident| fhir::PathSegment { ident: *ident, args: &[], bindings: &[] })
+                .chain([self.desugar_last_segment(res, *last, &path.generics)?]),
+        );
         let refine =
             try_alloc_slice!(self.genv(), &path.refine, |arg| self.desugar_refine_arg(arg))?;
-        Ok(fhir::Path { res, args, bindings, refine, span: path.span })
+        Ok(fhir::Path { res, segments, refine, span: path.span })
+    }
+
+    fn desugar_last_segment(
+        &mut self,
+        res: Res,
+        ident: surface::Ident,
+        args: &[surface::GenericArg],
+    ) -> Result<fhir::PathSegment<'genv>> {
+        let (args, bindings) = self.desugar_generic_args(res, args)?;
+        Ok(fhir::PathSegment { ident, args, bindings })
     }
 
     fn mk_lft_hole(&self) -> fhir::Lifetime {

crates/flux-fhir-analysis/src/annot_check.rs

@@ -348,11 +348,12 @@ impl<'zip, 'genv, 'tcx> Zipper<'zip, 'genv, 'tcx> {
             }
             return Err(self.emit_err(errors::InvalidRefinement::from_paths(path, expected_path)));
         }
-        if path.args.len() != expected_path.args.len() {
+        // TODO(nilehmann) we should check all segments here
+        if path.last_segment().args.len() != expected_path.last_segment().args.len() {
             return Err(self.emit_err(errors::GenericArgCountMismatch::new(path, expected_path)));
         }
 
-        iter::zip(path.args, expected_path.args)
+        iter::zip(path.last_segment().args, expected_path.last_segment().args)
             .try_for_each_exhaust(|(arg, expected)| self.zip_generic_arg(arg, expected))
     }
 
@@ -467,8 +468,8 @@ mod errors {
         pub(super) fn new(path: &fhir::Path, expected_path: &fhir::Path) -> Self {
             GenericArgCountMismatch {
                 span: path.span,
-                found: path.args.len(),
-                expected: expected_path.args.len(),
+                found: path.last_segment().args.len(),
+                expected: expected_path.last_segment().args.len(),
                 def_descr: path.res.descr(),
                 expected_span: expected_path.span,
             }

crates/flux-fhir-analysis/src/conv.rs

@@ -377,7 +377,7 @@ impl<'a, 'genv, 'tcx> ConvCtxt<'a, 'genv, 'tcx> {
         bounded_ty: &rty::Ty,
         bound: &fhir::GenericBound,
         clauses: &mut Vec<rty::Clause>,
-    ) -> QueryResult<()> {
+    ) -> QueryResult {
         match bound {
             fhir::GenericBound::Trait(trait_ref, fhir::TraitBoundModifier::None) => {
                 let trait_id = trait_ref.trait_def_id();
@@ -386,8 +386,21 @@ impl<'a, 'genv, 'tcx> ConvCtxt<'a, 'genv, 'tcx> {
                 } else {
                     let path = &trait_ref.trait_ref;
                     let params = &trait_ref.bound_generic_params;
-                    self.conv_trait_bound(env, bounded_ty, trait_id, path.args, params, clauses)?;
-                    self.conv_type_bindings(env, bounded_ty, trait_id, path.bindings, clauses)
+                    self.conv_trait_bound(
+                        env,
+                        bounded_ty,
+                        trait_id,
+                        path.last_segment().args,
+                        params,
+                        clauses,
+                    )?;
+                    self.conv_type_bindings(
+                        env,
+                        bounded_ty,
+                        trait_id,
+                        path.last_segment().bindings,
+                        clauses,
+                    )
                 }
             }
             // Maybe bounds are only supported for `?Sized`. The effect of the maybe bound is to
@@ -409,7 +422,7 @@ impl<'a, 'genv, 'tcx> ConvCtxt<'a, 'genv, 'tcx> {
         args: &[fhir::GenericArg],
         params: &[fhir::GenericParam],
         clauses: &mut Vec<rty::Clause>,
-    ) -> QueryResult<()> {
+    ) -> QueryResult {
         let mut into = vec![rty::GenericArg::Ty(bounded_ty.clone())];
         self.conv_generic_args_into(env, args, &mut into)?;
         self.fill_generic_args_defaults(trait_id, &mut into)?;
@@ -458,8 +471,8 @@ impl<'a, 'genv, 'tcx> ConvCtxt<'a, 'genv, 'tcx> {
 
         let pred = rty::FnTraitPredicate {
             self_ty: self_ty.clone(),
-            tupled_args: self.conv_ty(env, path.args[0].expect_type())?,
-            output: self.conv_ty(env, &path.bindings[0].term)?,
+            tupled_args: self.conv_ty(env, path.last_segment().args[0].expect_type())?,
+            output: self.conv_ty(env, &path.last_segment().bindings[0].term)?,
             kind,
         };
         // FIXME(nilehmann) We should use `tcx.late_bound_vars` here instead of trusting our lowering
@@ -725,8 +738,12 @@ impl<'a, 'genv, 'tcx> ConvCtxt<'a, 'genv, 'tcx> {
         if let fhir::BaseTyKind::Path(fhir::QPath::Resolved(self_ty, path)) = &bty.kind {
             if let fhir::Res::Def(DefKind::AssocTy, def_id) = path.res {
                 let self_ty = self.conv_ty(env, self_ty.as_deref().unwrap())?;
+                let [.., trait_segment, assoc_segment] = path.segments else {
+                    span_bug!(bty.span, "expected at least two segments");
+                };
                 let mut args = vec![rty::GenericArg::Ty(self_ty)];
-                args.append(&mut self.conv_generic_args(env, def_id, path.args)?);
+                self.conv_generic_args_into(env, trait_segment.args, &mut args)?;
+                self.conv_generic_args_into(env, assoc_segment.args, &mut args)?;
                 let args = List::from_vec(args);
 
                 let refine_args = List::empty();
@@ -866,7 +883,7 @@ impl<'a, 'genv, 'tcx> ConvCtxt<'a, 'genv, 'tcx> {
             }
             fhir::Res::Def(DefKind::Struct | DefKind::Enum, did) => {
                 let adt_def = self.genv.adt_def(*did)?;
-                let args = self.conv_generic_args(env, *did, path.args)?;
+                let args = self.conv_generic_args(env, *did, path.last_segment().args)?;
                 rty::BaseTy::adt(adt_def, args)
             }
             fhir::Res::Def(DefKind::TyParam, def_id) => {
@@ -881,7 +898,7 @@ impl<'a, 'genv, 'tcx> ConvCtxt<'a, 'genv, 'tcx> {
                     .replace_bound_expr(&idx));
             }
             fhir::Res::Def(DefKind::TyAlias { .. }, def_id) => {
-                let generics = self.conv_generic_args(env, *def_id, path.args)?;
+                let generics = self.conv_generic_args(env, *def_id, path.last_segment().args)?;
                 let refine = path
                     .refine
                     .iter()

crates/flux-fhir-analysis/src/wf/mod.rs

@@ -485,12 +485,14 @@ impl<'genv, 'tcx> Wf<'genv, 'tcx> {
         }
         let snapshot = self.xi.snapshot();
 
+        // TODO(nilehmann) we should check all segments
+        let last_segment = path.last_segment();
         if let fhir::Res::Def(_kind, did) = &path.res
-            && !path.args.is_empty()
+            && !last_segment.args.is_empty()
         {
-            self.check_generic_args(infcx, *did, path.args)?;
+            self.check_generic_args(infcx, *did, last_segment.args)?;
         }
-        let bindings = self.check_type_bindings(infcx, path.bindings);
+        let bindings = self.check_type_bindings(infcx, last_segment.bindings);
         if !self.genv.is_box(path.res) {
             self.xi.rollback_to(snapshot);
         }

crates/flux-middle/src/fhir.rs

@@ -604,12 +604,24 @@ pub enum QPath<'fhir> {
 #[derive(Clone, Copy)]
 pub struct Path<'fhir> {
     pub res: Res,
-    pub args: &'fhir [GenericArg<'fhir>],
-    pub bindings: &'fhir [TypeBinding<'fhir>],
+    pub segments: &'fhir [PathSegment<'fhir>],
     pub refine: &'fhir [RefineArg<'fhir>],
     pub span: Span,
 }
 
+impl<'fhir> Path<'fhir> {
+    pub fn last_segment(&self) -> &'fhir PathSegment<'fhir> {
+        self.segments.last().unwrap()
+    }
+}
+
+#[derive(Clone, Copy)]
+pub struct PathSegment<'fhir> {
+    pub ident: SurfaceIdent,
+    pub args: &'fhir [GenericArg<'fhir>],
+    pub bindings: &'fhir [TypeBinding<'fhir>],
+}
+
 #[derive(Clone, Copy)]
 pub struct TypeBinding<'fhir> {
     pub ident: SurfaceIdent,
@@ -1208,24 +1220,17 @@ impl fmt::Debug for QPath<'_> {
 
 impl fmt::Debug for Path<'_> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match self.res {
-            Res::PrimTy(PrimTy::Int(int_ty)) => {
-                write!(f, "{}", int_ty.name_str())?;
-            }
-            Res::PrimTy(PrimTy::Uint(uint_ty)) => {
-                write!(f, "{}", uint_ty.name_str())?;
-            }
-            Res::PrimTy(PrimTy::Float(float_ty)) => {
-                write!(f, "{}", float_ty.name_str())?;
-            }
-            Res::PrimTy(PrimTy::Bool) => write!(f, "bool")?,
-            Res::PrimTy(PrimTy::Str) => write!(f, "str")?,
-            Res::PrimTy(PrimTy::Char) => write!(f, "char")?,
-            Res::Def(_, def_id) => {
-                write!(f, "{}", pretty::def_id_to_string(def_id))?;
-            }
-            Res::SelfTyAlias { .. } | Res::SelfTyParam { .. } => write!(f, "Self")?,
+        write!(f, "{:?}", self.segments.iter().format("::"))?;
+        if !self.refine.is_empty() {
+            write!(f, "({:?})", self.refine.iter().format(", "))?;
         }
+        Ok(())
+    }
+}
+
+impl fmt::Debug for PathSegment<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.ident)?;
         let args: Vec<_> = self
             .args
             .iter()
@@ -1235,9 +1240,6 @@ impl fmt::Debug for Path<'_> {
         if !args.is_empty() {
             write!(f, "<{:?}>", args.iter().format(", "))?;
         }
-        if !self.refine.is_empty() {
-            write!(f, "({:?})", self.refine.iter().format(", "))?;
-        }
         Ok(())
     }
 }

crates/flux-middle/src/fhir/lift.rs

@@ -88,10 +88,14 @@ pub fn lift_self_ty<'genv>(
         let cx = LiftCtxt::new(genv, owner_id, &local_id_gen, None);
 
         let span = item.ident.span.to(generics.span);
-        let path = fhir::Path {
-            res: fhir::Res::Def(def_kind, owner_id.to_def_id()),
+        let segment = fhir::PathSegment {
+            ident: item.ident,
             args: cx.generic_params_into_args(generics)?,
             bindings: &[],
+        };
+        let path = fhir::Path {
+            res: fhir::Res::Def(def_kind, owner_id.to_def_id()),
+            segments: genv.alloc_slice(&[segment]),
             refine: &[],
             span,
         };
@@ -352,10 +356,10 @@ impl<'a, 'genv, 'tcx> LiftCtxt<'a, 'genv, 'tcx> {
         generics: &hir::Generics,
     ) -> fhir::VariantRet<'genv> {
         let span = item.ident.span.to(generics.span);
+        let segment = fhir::PathSegment { ident: item.ident, args: &[], bindings: &[] };
         let path = fhir::Path {
             res: fhir::Res::SelfTyAlias { alias_to: self.owner.to_def_id(), is_trait_impl: false },
-            args: &[],
-            bindings: &[],
+            segments: self.genv.alloc_slice(&[segment]),
             refine: &[],
             span,
         };
@@ -451,23 +455,25 @@ impl<'a, 'genv, 'tcx> LiftCtxt<'a, 'genv, 'tcx> {
         let Ok(res) = path.res.try_into() else {
             return self.emit_unsupported(&format!("unsupported res: `{:?}`", path.res));
         };
-        // TODO(RJ):REFACTOR-PATH Gross hack to account for <Vec<T> as Index<I>>::Output
-        let args = if let fhir::Res::Def(DefKind::AssocTy, _) = res {
-            path.segments.first().unwrap().args
-        } else {
-            path.segments.last().unwrap().args
-        };
+        let segments =
+            try_alloc_slice!(self.genv, path.segments, |segment| self.lift_path_segment(segment))?;
+
+        Ok(fhir::Path { res, segments, refine: &[], span: path.span })
+    }
 
+    fn lift_path_segment(
+        &mut self,
+        segment: &hir::PathSegment,
+    ) -> Result<fhir::PathSegment<'genv>> {
         let (args, bindings) = {
-            match args {
+            match segment.args {
                 Some(args) => {
                     (self.lift_generic_args(args.args)?, self.lift_type_bindings(args.bindings)?)
                 }
                 None => ([].as_slice(), [].as_slice()),
             }
         };
-
-        Ok(fhir::Path { res, args, bindings, refine: &[], span: path.span })
+        Ok(fhir::PathSegment { ident: segment.ident, args, bindings })
     }
 
     fn lift_path_to_ty(
@@ -552,8 +558,14 @@ impl<'a, 'genv, 'tcx> LiftCtxt<'a, 'genv, 'tcx> {
             match param.kind {
                 hir::GenericParamKind::Type { .. } => {
                     let res = fhir::Res::Def(DefKind::TyParam, param.def_id.to_def_id());
-                    let path =
-                        fhir::Path { res, args: &[], bindings: &[], refine: &[], span: param.span };
+                    let segment =
+                        fhir::PathSegment { ident: param.name.ident(), args: &[], bindings: &[] };
+                    let path = fhir::Path {
+                        res,
+                        segments: self.genv.alloc_slice(&[segment]),
+                        refine: &[],
+                        span: param.span,
+                    };
                     let bty = fhir::BaseTy::from(fhir::QPath::Resolved(None, path));
                     let ty = fhir::Ty { kind: fhir::TyKind::BaseTy(bty), span: param.span };
                     Ok(fhir::GenericArg::Type(self.genv.alloc(ty)))

crates/flux-middle/src/fhir/visit.rs

@@ -1,8 +1,8 @@
 use super::{
     AliasReft, BaseTy, BaseTyKind, Constraint, EnumDef, Expr, ExprKind, FieldDef, FnDecl, FnOutput,
-    FnSig, FuncSort, GenericArg, Generics, Ident, Lifetime, Lit, PolyFuncSort, QPath, RefineArg,
-    RefineArgKind, RefineParam, Sort, SortPath, StructDef, Ty, TyKind, TypeBinding, VariantDef,
-    VariantRet,
+    FnSig, FuncSort, GenericArg, Generics, Ident, Lifetime, Lit, Path, PathSegment, PolyFuncSort,
+    QPath, RefineArg, RefineArgKind, RefineParam, Sort, SortPath, StructDef, Ty, TyKind,
+    TypeBinding, VariantDef, VariantRet,
 };
 
 #[macro_export]
@@ -80,6 +80,14 @@ pub trait Visitor: Sized {
         walk_qpath(self, qpath);
     }
 
+    fn visit_path(&mut self, path: &Path) {
+        walk_path(self, path);
+    }
+
+    fn visit_path_segment(&mut self, segment: &PathSegment) {
+        walk_path_segment(self, segment);
+    }
+
     fn visit_type_binding(&mut self, binding: &TypeBinding) {
         walk_type_binding(self, binding);
     }
@@ -242,13 +250,21 @@ pub fn walk_qpath<V: Visitor>(vis: &mut V, qpath: &QPath) {
             if let Some(self_ty) = self_ty {
                 vis.visit_ty(self_ty);
             }
-            walk_list!(vis, visit_generic_arg, path.args);
-            walk_list!(vis, visit_type_binding, path.bindings);
             walk_list!(vis, visit_refine_arg, path.refine);
         }
     }
 }
 
+pub fn walk_path<V: Visitor>(vis: &mut V, path: &Path) {
+    walk_list!(vis, visit_path_segment, path.segments);
+    walk_list!(vis, visit_refine_arg, path.refine);
+}
+
+pub fn walk_path_segment<V: Visitor>(vis: &mut V, segment: &PathSegment) {
+    walk_list!(vis, visit_generic_arg, segment.args);
+    walk_list!(vis, visit_type_binding, segment.bindings);
+}
+
 pub fn walk_type_binding<V: Visitor>(vis: &mut V, binding: &TypeBinding) {
     let TypeBinding { ident: _, term } = binding;
     vis.visit_ty(term);

crates/flux-middle/src/sort_of.rs

@@ -32,7 +32,8 @@ impl<'sess, 'tcx> GlobalEnv<'sess, 'tcx> {
             fhir::Res::Def(DefKind::TyAlias { .. } | DefKind::Enum | DefKind::Struct, def_id) => {
                 let mut sort_args = vec![];
                 let sort_def = self.adt_sort_def_of(def_id);
-                for arg in sort_def.filter_generic_args(path.args) {
+                let generic_args = path.segments.last().unwrap().args;
+                for arg in sort_def.filter_generic_args(generic_args) {
                     sort_args.push(self.sort_of_ty(arg.expect_type())?);
                 }
                 let ctor = rty::SortCtor::Adt(self.adt_sort_def_of(def_id));