diff --git a/src/back/hlsl/writer.rs b/src/back/hlsl/writer.rs
index 6da98502f1..e9809bd2d4 100644
--- a/src/back/hlsl/writer.rs
+++ b/src/back/hlsl/writer.rs
@@ -2352,12 +2352,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 }
             }
             Expression::FunctionArgument(pos) => {
-                let key = match func_ctx.ty {
-                    back::FunctionType::Function(handle) => NameKey::FunctionArgument(handle, pos),
-                    back::FunctionType::EntryPoint(index) => {
-                        NameKey::EntryPointArgument(index, pos)
-                    }
-                };
+                let key = func_ctx.argument_key(pos);
                 let name = &self.names[&key];
                 write!(self.out, "{name}")?;
             }
diff --git a/src/back/mod.rs b/src/back/mod.rs
index 875b04d69a..1314e14d77 100644
--- a/src/back/mod.rs
+++ b/src/back/mod.rs
@@ -37,13 +37,28 @@ impl std::fmt::Display for Level {
     }
 }
 
-/// Stores the current function type (either a regular function or an entry point)
+/// Whether we're generating an entry point or a regular function.
 ///
-/// Also stores data needed to identify it (handle for a regular function or index for an entry point)
+/// Backend languages often require different code for a [`Function`]
+/// depending on whether it represents an [`EntryPoint`] or not.
+/// Backends can pass common code one of these values to select the
+/// right behavior.
+///
+/// These values also carry enough information to find the `Function`
+/// in the [`Module`]: the `Handle` for a regular function, or the
+/// index into [`Module::entry_points`] for an entry point.
+///
+/// [`Function`]: crate::Function
+/// [`EntryPoint`]: crate::EntryPoint
+/// [`Module`]: crate::Module
+/// [`Module::entry_points`]: crate::Module::entry_points
 enum FunctionType {
-    /// A regular function and it's handle
+    /// A regular function.
     Function(crate::Handle<crate::Function>),
-    /// A entry point and it's index
+    /// An [`EntryPoint`], and its index in [`Module::entry_points`].
+    ///
+    /// [`EntryPoint`]: crate::EntryPoint
+    /// [`Module::entry_points`]: crate::Module::entry_points
     EntryPoint(crate::proc::EntryPointIndex),
 }
 
diff --git a/src/back/wgsl/writer.rs b/src/back/wgsl/writer.rs
index 3a0424e7bb..980d768782 100644
--- a/src/back/wgsl/writer.rs
+++ b/src/back/wgsl/writer.rs
@@ -247,14 +247,7 @@ impl<W: Write> Writer<W> {
                 self.write_attributes(&map_binding_to_attribute(binding))?;
             }
             // Write argument name
-            let argument_name = match func_ctx.ty {
-                back::FunctionType::Function(handle) => {
-                    &self.names[&NameKey::FunctionArgument(handle, index as u32)]
-                }
-                back::FunctionType::EntryPoint(ep_index) => {
-                    &self.names[&NameKey::EntryPointArgument(ep_index, index as u32)]
-                }
-            };
+            let argument_name = &self.names[&func_ctx.argument_key(index as u32)];
 
             write!(self.out, "{argument_name}: ")?;
             // Write argument type