Re: [PATCH v2 1/2] rust: add initial scatterlist bindings

["Alexandre Courbot" <acourbot@nvidia.com>]

On Mon Jun 30, 2025 at 9:56 PM JST, Alexandre Courbot wrote:
> On Mon Jun 30, 2025 at 5:34 PM JST, Alexandre Courbot wrote:
>> I actually have some more comments, but I'd like to understasnd first
>> why you decided to drop the typestate. If there is no good reason, I
>> think I can provide a more detailed explanation about the design I had
>> in mind when thinking about Lyude's usecase - basically, a two-stages
>> typestate SGTable type where the first stage is unsafe, but the second
>> one leverages SGTablePages and is safe. But going into that explanation
>> now would be pointless if the typestate cannot be used for some reason.
>
> After thinking some more about it, I think I can verbalize better my
> expectations for this interface (and my problems with this version):
>
> Basically, I believe we should (and can) only need unsafe methods when
> using one of the constructors for a SG table. Once the SG table object
> is built, nothing in the interface should need to be unsafe, and we
> should have access to exactly the features that are safe to use
> according to the construction-time invariants. This implies that we have
> several constructors, and several types for SG tables - possibly
> typestates or completely distinct types as you did in this version.
>
> I wrote a bit of test code trying to match both the requirements of GEM
> SHMEM objects (work with an already-existing `sg_table`), and of owned
> SG tables ; and I *think* I got something that works by leveraging
> `Borrow`. Let me clean up my code after a good night of sleep and I'll
> try to elaborate a bit.

Alright, that was an interesting rabbit hole, but I think we're closing in.
Apologies, for this is a long post.

Let's first summarize the use-cases we want to support:

- The GEM use-case, where we need to work with `sg_tables` that already exist
  and are owned by some existing entity. We want to "wrap" these into a Rust
  type that provides methods that are always safe to call. In this case the
  Rust `SGTable` does not allocate or free the underlying `sg_table` ; thus
  it is primordial to ensure that it 1) doesn't outlive the `sg_table` and
  2) that the `sg_table` is not modified while we are using it. Let's call it
  the "borrowed" case.

- The Nova use-case, where we want to make an bit of memory (e.g. a `VVec`)
  available to a device. In this case, the Rust `SGTable` DOES own the
  `sg_table` and can even mutate it. However, it shall not outlive the
  backing memory, which must be pinned for as long as the `SGTable` exists.
  Let's call it the "owned" case.

For the borrowed case, there is not much choice but to use an `unsafe`
constructor, with the caller guaranteeing that the expected properties of the
`sg_table` are respected. Conversely, the owned case can be built safely for
any type that provides an implementation of the `SGTablePages` trait (we need
to find a more explicit name for this trait).

So, with this in mind, there are two dimensions around which a `SGTable` can be
built:

1. The ownership or not of the underlying `sg_table`,
2. Whether the `sg_table` is DMA-mapped or not.

For 1., the `Borrow` and `BorrowMut` traits have us covered. If we access the
`sg_table` through them, we can support both the borrowed and owned scenarios.

For 2., a typestate can ensure that only methods that are valid for the given
mapped state of the `SGTable` are visible.

Which makes our `SGTable` look something like this:

  pub struct SGTable<T: Borrow<bindings::sg_table>, M: MappingState> {
      // Declared first so it is dropped first, as we want to unmap before
      // freeing the `sg_table` if we own it.
      mapping: M,
      table: T,
  }

With the mapping typestate being:

  pub trait MappingState {}

  // For sg_tables that are not mapped.
  pub struct Unmapped;
  impl MappingState for Unmapped {}

  // For sg_tables that are mapped, but not managed by us.
  pub struct BorrowedMapping;
  impl MappingState for BorrowedMapping {}

This lets us have constructors to cover the case where we want to wrap an
existing `sg_table`:

  impl<T> SGTable<T, Unmapped>
  where
      T: Borrow<bindings::sg_table>,
  {
      // Safety: 'r' must borrow a sg_table that is unmapped, and which
      // does not change as long as the returned object exists.
      pub unsafe fn new_unmapped(r: T) -> Self {
          Self {
              table: r,
              mapping: Unmapped,
          }
      }
  }

  impl<T> SGTable<T, BorrowedMapping>
  where
      T: Borrow<bindings::sg_table>,
  {
      // Safety: 'r' must borrow a sg_table that is DMA-mapped, and which
      // does not change as long as the returned object exists.
      pub unsafe fn new_mapped(r: T) -> Self {
          Self {
              table: r,
              mapping: BorrowedMapping,
          }
      }
  }

And this should be enough to cover the needs of GEM. Lyude mentioned building a
wrapper around a reference to a `sg_table`, which can be done as follows:

  // Obtain the reference from `gem_shmem_object` with the proper lifetime.
  let sg_table_ref: &bindings::sg_table = ...
  let sg_table = SGTable::new_mapped(sg_table_ref);

Here `SGTable` borrows a reference to `gem_shmem_object`, meaning it cannot
outlive it, and `gem_shmem_object` cannot be mutated for as long as we hold
that reference.

This also works to implement an equivalent of `OwnedSGTable`, if I understood
it correctly:

  struct WrappedAref(ARef<gem::Object>);
  impl Borrow<bindings::sg_table> for WrapperARef ...

  // `self` is a `&gem::Object`.
  let wrapped_aref: WrapperAref = self.into();
  let sg_table = SGTable::new_mapped(wrapped_aref);

Here the fact that we are passing an `ARef` lets the GEM object's lifetime be
managed at runtime, just like `OwnedSGTable` does.

Now on to the Nova use-case. Here we want to allocate, manage, and eventually
free both the `struct sg_table` and its DMA mapping.

For the former, we can define a new struct that implements `Borrow` and takes
care of freeing the `sg_table`:

  pub struct OwnedSgt<P: SGTablePages> {
        sgt: bindings::sg_table,
        pages: P,
  }

  impl<P> Drop for OwnedSgt<P>
    where
        P: SGTablePages,
  {
        fn drop(&mut self) {
          unsafe { bindings::sg_free_table(&mut self.sgt) };
      }
  }

  impl<P> Borrow<bindings::sg_table> for OwnedSgt<P>
    where
        P: SGTablePages,
  {
        fn borrow(&self) -> &bindings::sg_table {
          &self.sgt
      }
  }

This will be our first generic argument for `SGTable`. The mapping can then be
handled similarly:

  pub struct ManagedMapping {
      dev: ARef<Device>,
      dir: DmaDataDirection,
      // This works because the `sgl` member of `struct sg_table` never moves
      // after being allocated.
      sgl: *mut bindings::scatterlist,
      orig_nents: ffi::c_uint,
  }
  impl MappingState for ManagedMapping {}

  impl Drop for ManagedMapping {
      fn drop(&mut self) {
          unsafe {
              bindings::dma_unmap_sg_attrs(
                  self.dev.as_raw(),
                  self.sgl,
                  self.orig_nents as i32,
                  self.dir as i32,
                  0,
              )
          };
      }
  }

With this, and the `SGTablePages` trait, we can now provide a safe constructor
for mapping existing memory into a device:

  impl<P: SGTablePages> SGTable<OwnedSgt<P>, ManagedMapping> {
        pub fn new_owned(pages: P, flags: kernel::alloc::Flags) -> Result<Self> ...

The DMA mapping then remains in effect for as long as the returned object
lives.

You would then have `impl` blocks to provide the raw `sg_table` pointer as well
as DMA or CPU address iterators, made available or not depending on the mapping
typestate. And if the type borrowing the `sg_table` also implements
`BorrowMut`, we can even change the mapping state programmatically. I have
omitted it here for conciseness but have some code for this as well.

Although I remember a mention of the `Unmapped` state being unneeded in
discussion of the previous iteration - and indeed, so far both the GEM and Nova
use-cases do not really need it, so if that's confirmed we could remove the
`Unmapped` state and any kind of transition to simplify the interface.

Thoughts? If Abdiel is comfortable with this I can submit a v3 with this design
for review (putting myself as co-developer), on which Abdiel could then keep
iterating, as I suspect this would be easier to understand than this long email
:).