nak: memstream: move into common code

Move the memstream code into common code. Other Rust code interfacing
with FILE pointers will find the memstream abstraction useful.

Most notably, pinning is actually enforced this time with PhantomPinned.

Add a .clang-format from a sibling dir (i.e.: compiler/nir) while we're
at it to keep things tidy.

Signed-off-by: Daniel Almeida <daniel.almeida@collabora.com>
Reviewed-by: Faith Ekstrand <faith.ekstrand@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30594>

src/compiler/rust/.clang-format

@@ -0,0 +1,6 @@
+BasedOnStyle: InheritParentConfig
+DisableFormat: false
+
+ColumnLimit: 0
+Cpp11BracedListStyle: false
+SpaceBeforeParens: ControlStatementsExceptControlMacros

src/compiler/rust/bindings.h

@@ -3,4 +3,6 @@
  * SPDX-License-Identifier: MIT
  */
 
+#include "util/memstream.h"
+#include "rust_helpers.h"
 #include "nir.h"

src/compiler/rust/lib.rs

@@ -5,5 +5,6 @@ pub mod as_slice;
 pub mod bindings;
 pub mod bitset;
 pub mod cfg;
+pub mod memstream;
 pub mod nir;
 pub mod smallvec;

src/compiler/rust/memstream.rs

@@ -0,0 +1,150 @@
+// Copyright © 2024 Collabora, Ltd.
+// SPDX-License-Identifier: MIT
+
+use std::io;
+use std::marker::PhantomPinned;
+use std::pin::Pin;
+
+use crate::bindings;
+
+struct MemStreamImpl {
+    stream: bindings::u_memstream,
+    buffer: *mut u8,
+    buffer_size: usize,
+    _pin: PhantomPinned,
+}
+
+/// A Rust memstream abstraction. Useful when interacting with C code that
+/// expects a FILE* pointer.
+///
+/// The size of the buffer is managed by the C code automatically.
+pub struct MemStream(Pin<Box<MemStreamImpl>>);
+
+impl MemStream {
+    pub fn new() -> io::Result<Self> {
+        let mut stream_impl = Box::pin(MemStreamImpl {
+            stream: unsafe { std::mem::zeroed() },
+            buffer: std::ptr::null_mut(),
+            buffer_size: 0,
+            _pin: PhantomPinned,
+        });
+
+        unsafe {
+            let stream_impl = stream_impl.as_mut().get_unchecked_mut();
+            if !bindings::u_memstream_open(
+                &mut stream_impl.stream,
+                (&mut stream_impl.buffer).cast(),
+                &mut stream_impl.buffer_size,
+            ) {
+                return Err(io::Error::last_os_error());
+            }
+            if bindings::u_memstream_flush(&mut stream_impl.stream) != 0 {
+                return Err(io::Error::last_os_error());
+            }
+        }
+
+        Ok(Self(stream_impl))
+    }
+
+    // Safety: caller must ensure that inner is not moved through the returned
+    // reference.
+    unsafe fn inner_mut(&mut self) -> &mut MemStreamImpl {
+        unsafe { self.0.as_mut().get_unchecked_mut() }
+    }
+
+    /// Flushes the stream so written data appears in the stream
+    pub fn flush(&mut self) -> io::Result<()> {
+        unsafe {
+            let stream = self.inner_mut();
+            if bindings::u_memstream_flush(&mut stream.stream) != 0 {
+                return Err(io::Error::last_os_error());
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Resets the MemStream
+    pub fn reset(&mut self) -> io::Result<()> {
+        *self = Self::new()?;
+        Ok(())
+    }
+
+    /// Resets the MemStream and returns its contents
+    pub fn take(&mut self) -> io::Result<Vec<u8>> {
+        let mut vec = Vec::new();
+        vec.extend_from_slice(self.as_slice()?);
+        self.reset()?;
+        Ok(vec)
+    }
+
+    /// Resets the MemStream and returns its contents as a UTF-8 string
+    pub fn take_utf8_string_lossy(&mut self) -> io::Result<String> {
+        let string = String::from_utf8_lossy(self.as_slice()?).into_owned();
+        self.reset()?;
+        Ok(string)
+    }
+
+    /// Returns the current position in the stream.
+    pub fn position(&self) -> usize {
+        unsafe { bindings::compiler_rs_ftell(self.c_file()) as usize }
+    }
+
+    /// Seek to a position relative to the start of the stream.
+    pub fn seek(&mut self, offset: u64) -> io::Result<()> {
+        let offset = offset.try_into().map_err(|_| {
+            io::Error::new(io::ErrorKind::InvalidInput, "offset too large")
+        })?;
+
+        unsafe {
+            if bindings::compiler_rs_fseek(self.c_file(), offset, 0) != 0 {
+                Err(io::Error::last_os_error())
+            } else {
+                Ok(())
+            }
+        }
+    }
+
+    /// Returns the underlying C file.
+    ///
+    /// # Safety
+    ///
+    /// The memstream abstraction assumes that the file is valid throughout its
+    /// lifetime.
+    pub unsafe fn c_file(&self) -> *mut bindings::FILE {
+        self.0.stream.f
+    }
+
+    /// Returns a slice view into the memstream
+    ///
+    /// This is only safe with respect to other safe Rust methods.  Even though
+    /// this takes a reference to the stream there is nothing preventing you
+    /// from modifying the stream through the FILE with unsafe C code.
+    ///
+    /// This is conceptually the same as `AsRef`, but it flushes the stream
+    /// first, which means it takes &mut self as a receiver.
+    fn as_slice(&mut self) -> io::Result<&[u8]> {
+        // Make sure we have the most up-to-date data before returning a slice.
+        self.flush()?;
+        let pos = self.position();
+
+        if pos == 0 {
+            Ok(&[])
+        } else {
+            // SAFETY: this does not move the stream and we know that
+            // self.position() cannot exceed the stream size as per the
+            // open_memstream() API.
+            Ok(unsafe { std::slice::from_raw_parts(self.0.buffer, pos) })
+        }
+    }
+}
+
+impl Drop for MemStream {
+    fn drop(&mut self) {
+        // SAFETY: this does not move the stream.
+        unsafe {
+            bindings::u_memstream_close(&mut self.inner_mut().stream);
+            bindings::compiler_rs_free(self.0.buffer as *mut std::ffi::c_void);
+        }
+    }
+}

src/compiler/rust/meson.build

@@ -5,6 +5,7 @@ _compiler_rs_sources = [
   'as_slice.rs',
   'bitset.rs',
   'cfg.rs',
+  'memstream.rs',
   'nir.rs',
   'smallvec.rs',
 ]
@@ -52,9 +53,13 @@ _compiler_bindgen_args = [
   '--raw-line', '#![allow(non_snake_case)]',
   '--raw-line', '#![allow(non_upper_case_globals)]',
   '--allowlist-var', 'nir_.*_infos',
+  '--allowlist-var', 'rust_.*',
   '--allowlist-function', 'glsl_.*',
   '--allowlist-function', '_mesa_shader_stage_to_string',
   '--allowlist-function', 'nir_.*',
+  '--allowlist-function', 'compiler_rs.*',
+  '--allowlist-function', 'u_memstream.*',
+  '--allowlist-type', 'u_memstream',
   '--no-prepend-enum-name',
 ]
 
@@ -62,6 +67,21 @@ foreach type : _compiler_binding_types
   _compiler_bindgen_args += ['--allowlist-type', type]
 endforeach
 
+_libcompiler_c_sources = files('rust_helpers.c')
+
+_libcompiler_c = static_library(
+  'compiler_c_helpers',
+  [_libcompiler_c_sources],
+  include_directories : [inc_include, inc_util],
+  c_args : [no_override_init_args],
+  gnu_symbol_visibility : 'hidden',
+)
+
+_idep_libcompiler_c = declare_dependency(
+  include_directories: include_directories('.'),
+  link_with : _libcompiler_c,
+)
+
 _compiler_bindings_rs = rust.bindgen(
   input : ['bindings.h'],
   output : 'bindings.rs',
@@ -71,6 +91,7 @@ _compiler_bindings_rs = rust.bindgen(
   args : _compiler_bindgen_args,
   dependencies : [
     idep_nir_headers,
+    idep_mesautil,
   ],
 )
 
@@ -91,6 +112,7 @@ _libcompiler_rs = static_library(
   _compiler_rs_sources,
   gnu_symbol_visibility : 'hidden',
   rust_abi : 'rust',
+  dependencies: [_idep_libcompiler_c],
 )
 
 idep_compiler_rs = declare_dependency(

src/compiler/rust/rust_helpers.c

@@ -0,0 +1,28 @@
+/*
+ * Copyright © 2024 Collabora, Ltd.
+ * SPDX-License-Identifier: MIT
+ *
+ * This file contains helpers that are implemented in C so that, among other
+ * things, we avoid pulling in all of libc as bindings only to access a few
+ * functions.
+ */
+
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "rust_helpers.h"
+
+void compiler_rs_free(void *ptr)
+{
+   free(ptr);
+}
+
+long compiler_rs_ftell(FILE *f)
+{
+   return ftell(f);
+}
+
+int compiler_rs_fseek(FILE *f, long offset, int whence)
+{
+   return fseek(f, offset, whence);
+}

src/compiler/rust/rust_helpers.h

@@ -0,0 +1,10 @@
+/*
+ * Copyright © 2024 Collabora, Ltd.
+ * SPDX-License-Identifier: MIT
+ */
+
+#include <stdio.h>
+
+void compiler_rs_free(void *ptr);
+long compiler_rs_ftell(FILE *f);
+int compiler_rs_fseek(FILE *f, long offset, int whence);

src/nouveau/compiler/meson.build

@@ -33,7 +33,6 @@ libnak_c_files = files(
   'nak_nir_lower_tex.c',
   'nak_nir_lower_vtg_io.c',
   'nak_nir_split_64bit_conversions.c',
-  'nak_memstream.c',
 )
 
 _libacorn_rs = static_library(

src/nouveau/compiler/nak/from_nir.rs

@@ -335,7 +335,7 @@ impl<'a> ShaderFromNir<'a> {
             end_block_id: 0,
             ssa_map: HashMap::new(),
             saturated: HashSet::new(),
-            nir_instr_printer: NirInstrPrinter::new(),
+            nir_instr_printer: NirInstrPrinter::new().unwrap(),
         }
     }
 
@@ -3302,6 +3302,7 @@ impl<'a> ShaderFromNir<'a> {
                 let annotation = self
                     .nir_instr_printer
                     .instr_to_string(ni)
+                    .unwrap()
                     .split_whitespace()
                     .collect::<Vec<_>>()
                     .join(" ");
@@ -3350,6 +3351,7 @@ impl<'a> ShaderFromNir<'a> {
                 let annotation = self
                     .nir_instr_printer
                     .instr_to_string(ni)
+                    .unwrap()
                     .split_whitespace()
                     .collect::<Vec<_>>()
                     .join(" ");
@@ -3427,6 +3429,7 @@ impl<'a> ShaderFromNir<'a> {
                     let annotation = self
                         .nir_instr_printer
                         .instr_to_string(ni)
+                        .unwrap()
                         .split_whitespace()
                         .collect::<Vec<_>>()
                         .join(" ");

src/nouveau/compiler/nak/nir_instr_printer.rs

@@ -1,53 +1,27 @@
 // Copyright © 2024 Collabora, Ltd.
 // SPDX-License-Identifier: MIT
 
-use std::pin::Pin;
+use std::io;
 
+use compiler::bindings;
 use compiler::bindings::nir_instr;
-use nak_bindings::nak_clear_memstream;
-use nak_bindings::nak_close_memstream;
-use nak_bindings::nak_memstream;
-use nak_bindings::nak_nir_asprint_instr;
-use nak_bindings::nak_open_memstream;
+use compiler::memstream::MemStream;
 
 /// A memstream that holds the printed NIR instructions.
 pub struct NirInstrPrinter {
-    //  XXX: we need this to be pinned because we've passed references to its
-    //  fields when calling open_memstream.
-    stream: Pin<Box<nak_memstream>>,
+    stream: MemStream,
 }
 
 impl NirInstrPrinter {
-    pub fn new() -> Self {
-        let mut stream =
-            Box::pin(unsafe { std::mem::zeroed::<nak_memstream>() });
-        unsafe {
-            nak_open_memstream(stream.as_mut().get_unchecked_mut());
-        }
-        Self { stream }
+    pub fn new() -> io::Result<Self> {
+        Ok(Self {
+            stream: MemStream::new()?,
+        })
     }
 
     /// Prints the given NIR instruction.
-    pub fn instr_to_string(&mut self, instr: &nir_instr) -> String {
-        unsafe {
-            let stream = self.stream.as_mut().get_unchecked_mut();
-            nak_nir_asprint_instr(stream, instr);
-            let bytes = std::slice::from_raw_parts(
-                stream.buffer as *const u8,
-                stream.written,
-            );
-            let string = String::from_utf8_lossy(bytes).into_owned();
-            nak_clear_memstream(stream);
-            string
-        }
-    }
-}
-
-impl Drop for NirInstrPrinter {
-    fn drop(&mut self) {
-        unsafe {
-            let stream = self.stream.as_mut().get_unchecked_mut();
-            nak_close_memstream(stream)
-        }
+    pub fn instr_to_string(&mut self, instr: &nir_instr) -> io::Result<String> {
+        unsafe { bindings::nir_print_instr(instr, self.stream.c_file()) };
+        self.stream.take_utf8_string_lossy()
     }
 }

src/nouveau/compiler/nak_memstream.c

@@ -1,39 +0,0 @@
-/*
- * Copyright © 2024 Collabora, Ltd.
- * SPDX-License-Identifier: MIT
- *
- * This file exposes a nice interface that can be consumed from Rust. We would
- * have to have Rust libc bindings otherwise.
- */
-
-#include <assert.h>
-#include <stdio.h>
-#include <stdlib.h>
-
-#include "nak_private.h"
-#include "nir.h"
-
-void nak_open_memstream(struct nak_memstream *memstream)
-{
-    memstream->stream = open_memstream(&memstream->buffer, &memstream->written);
-    fflush(memstream->stream);
-    assert(memstream->stream);
-    assert(memstream->buffer);
-}
-
-void nak_close_memstream(struct nak_memstream *memstream)
-{
-    fclose(memstream->stream);
-    free(memstream->buffer);
-}
-
-void nak_nir_asprint_instr(struct nak_memstream *memstream, const nir_instr *instr)
-{
-    nir_print_instr(instr, memstream->stream);
-    fflush(memstream->stream);
-}
-
-void nak_clear_memstream(struct nak_memstream *memstream)
-{
-    rewind(memstream->stream);
-}

src/nouveau/compiler/nak_private.h

@@ -233,18 +233,6 @@ bool nak_nir_lower_cf(nir_shader *nir);
 
 void nak_optimize_nir(nir_shader *nir, const struct nak_compiler *nak);
 
-struct nak_memstream {
-   FILE *stream;
-   char *buffer;
-   size_t written;
-};
-
-void nak_open_memstream(struct nak_memstream *memstream);
-void nak_close_memstream(struct nak_memstream *memstream);
-void nak_flush_memstream(struct nak_memstream *memstream);
-void nak_clear_memstream(struct nak_memstream *memstream);
-void nak_nir_asprint_instr(struct nak_memstream *memstream, const nir_instr *instr);
-
 #ifdef __cplusplus
 }
 #endif