From 1421319dcfbaf79f76552f024b7e39168b73ea45 Mon Sep 17 00:00:00 2001
From: Christian Gmeiner <cgmeiner@igalia.com>
Date: Thu, 26 Sep 2024 22:34:45 +0200
Subject: [PATCH] compiler/rust: Copy MappedInstrs from NAK

Rename it to SmallVec, make it more generic and switch NAK
to it.

Signed-off-by: Christian Gmeiner <cgmeiner@igalia.com>
Reviewed-by: Faith Ekstrand <faith.ekstrand@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/31409>
---
 src/compiler/rust/lib.rs       |  1 +
 src/compiler/rust/meson.build  |  1 +
 src/compiler/rust/smallvec.rs  | 83 ++++++++++++++++++++++++++++++++++
 src/nouveau/compiler/nak/ir.rs | 37 +--------------
 4 files changed, 87 insertions(+), 35 deletions(-)
 create mode 100644 src/compiler/rust/smallvec.rs

diff --git a/src/compiler/rust/lib.rs b/src/compiler/rust/lib.rs
index 567db097ac6..676805cdecc 100644
--- a/src/compiler/rust/lib.rs
+++ b/src/compiler/rust/lib.rs
@@ -6,3 +6,4 @@ pub mod bindings;
 pub mod bitset;
 pub mod cfg;
 pub mod nir;
+pub mod smallvec;
diff --git a/src/compiler/rust/meson.build b/src/compiler/rust/meson.build
index c6928580888..4b5278489dd 100644
--- a/src/compiler/rust/meson.build
+++ b/src/compiler/rust/meson.build
@@ -6,6 +6,7 @@ _compiler_rs_sources = [
   'bitset.rs',
   'cfg.rs',
   'nir.rs',
+  'smallvec.rs',
 ]
 
 bindgen_version = find_program('bindgen').version()
diff --git a/src/compiler/rust/smallvec.rs b/src/compiler/rust/smallvec.rs
new file mode 100644
index 00000000000..099eb792442
--- /dev/null
+++ b/src/compiler/rust/smallvec.rs
@@ -0,0 +1,83 @@
+// Copyright © 2022 Collabora, Ltd.
+// SPDX-License-Identifier: MIT
+
+/// `SmallVec` is an optimized data structure that handles collections of items.
+/// It is designed to avoid allocating a `Vec` unless multiple items are present.
+///
+/// # Variants
+///
+/// * `None` - Represents an empty collection, no items are stored.
+/// * `One(T)` - Stores a single item without allocating a `Vec`.
+/// * `Many(Vec<T>)` - Stores multiple items in a heap-allocated `Vec`.
+///
+/// This helps to reduce the amount of Vec's allocated in the optimization passes.
+pub enum SmallVec<T> {
+    None,
+    One(T),
+    Many(Vec<T>),
+}
+
+impl<T> SmallVec<T> {
+    /// Adds an item to the `SmallVec`.
+    ///
+    /// If the collection is empty (`None`), the item is stored as `One`.
+    /// If the collection has one item (`One`), it transitions to `Many` and both items are stored in a `Vec`.
+    /// If the collection is already in the `Many` variant, the new item is pushed into the existing `Vec`.
+    ///
+    /// # Arguments
+    ///
+    /// * `item` - The item to be added.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// let mut vec: SmallVec<String> = SmallVec::None;
+    /// vec.push("Hello".to_string());
+    /// vec.push("World".to_string());
+    /// ```
+    pub fn push(&mut self, i: T) {
+        match self {
+            SmallVec::None => {
+                *self = SmallVec::One(i);
+            }
+            SmallVec::One(_) => {
+                *self = match std::mem::replace(self, SmallVec::None) {
+                    SmallVec::One(o) => SmallVec::Many(vec![o, i]),
+                    _ => panic!("Not a One"),
+                };
+            }
+            SmallVec::Many(v) => {
+                v.push(i);
+            }
+        }
+    }
+
+    /// Returns a mutable reference to the last item in the `SmallVec`, if it exists.
+    ///
+    /// * If the collection is empty (`None`), it returns `None`.
+    /// * If the collection has one item (`One`), it returns a mutable reference to that item.
+    /// * If the collection has multiple items (`Many`), it returns a mutable reference to the last item in the `Vec`.
+    ///
+    /// # Returns
+    ///
+    /// * `Option<&mut T>` - A mutable reference to the last item, or `None` if the collection is empty.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// let mut vec: SmallVec<i32> = SmallVec::None;
+    /// vec.push(1);
+    /// vec.push(2);
+    ///
+    /// if let Some(last) = vec.last_mut() {
+    ///     *last = 10;  // Modify the last element.
+    /// }
+    /// ```
+    pub fn last_mut(&mut self) -> Option<&mut T> {
+        match self {
+            SmallVec::None => None,
+            SmallVec::One(item) => Some(item),
+            SmallVec::Many(v) => v.last_mut(),
+        }
+    }
+}
diff --git a/src/nouveau/compiler/nak/ir.rs b/src/nouveau/compiler/nak/ir.rs
index c268bcd87cf..66b6a213f0f 100644
--- a/src/nouveau/compiler/nak/ir.rs
+++ b/src/nouveau/compiler/nak/ir.rs
@@ -12,6 +12,7 @@ use crate::legalize::LegalizeBuilder;
 use crate::sph::{OutputTopology, PixelImap};
 use compiler::as_slice::*;
 use compiler::cfg::CFG;
+use compiler::smallvec::SmallVec;
 use nak_ir_proc::*;
 use std::cmp::{max, min};
 use std::fmt;
@@ -6783,41 +6784,7 @@ impl<T: Into<Op>> From<T> for Instr {
     }
 }
 
-/// The result of map() done on a Box<Instr>. A Vec is only allocated if the
-/// mapping results in multiple instructions. This helps to reduce the amount of
-/// Vec's allocated in the optimization passes.
-pub enum MappedInstrs {
-    None,
-    One(Box<Instr>),
-    Many(Vec<Box<Instr>>),
-}
-
-impl MappedInstrs {
-    pub fn push(&mut self, i: Box<Instr>) {
-        match self {
-            MappedInstrs::None => {
-                *self = MappedInstrs::One(i);
-            }
-            MappedInstrs::One(_) => {
-                *self = match std::mem::replace(self, MappedInstrs::None) {
-                    MappedInstrs::One(o) => MappedInstrs::Many(vec![o, i]),
-                    _ => panic!("Not a One"),
-                };
-            }
-            MappedInstrs::Many(v) => {
-                v.push(i);
-            }
-        }
-    }
-
-    pub fn last_mut(&mut self) -> Option<&mut Box<Instr>> {
-        match self {
-            MappedInstrs::None => None,
-            MappedInstrs::One(instr) => Some(instr),
-            MappedInstrs::Many(v) => v.last_mut(),
-        }
-    }
-}
+pub type MappedInstrs = SmallVec<Box<Instr>>;
 
 pub struct BasicBlock {
     pub label: Label,