From 1355c7194302bc6e0b37d03afd727cc6c52c1994 Mon Sep 17 00:00:00 2001
From: Faith Ekstrand <faith.ekstrand@collabora.com>
Date: Wed, 12 Mar 2025 16:50:49 -0500
Subject: [PATCH] compiler/rust: Add a nir_alu_type wrapper

Reviewed-by: Mel Henning <mhenning@darkrefraction.com>
Reviewed-by: Benjamin Lee <benjamin.lee@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/34126>
---
 src/compiler/rust/meson.build |  1 +
 src/compiler/rust/nir.rs      | 31 +++++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/src/compiler/rust/meson.build b/src/compiler/rust/meson.build
index 62835ebd8c9..c16d72f157a 100644
--- a/src/compiler/rust/meson.build
+++ b/src/compiler/rust/meson.build
@@ -39,6 +39,7 @@ _compiler_bindgen_args = [
   '--raw-line', '#![allow(non_camel_case_types)]',
   '--raw-line', '#![allow(non_snake_case)]',
   '--raw-line', '#![allow(non_upper_case_globals)]',
+  '--allowlist-var', 'NIR_.*',
   '--allowlist-var', 'nir_.*_infos',
   '--allowlist-var', 'rust_.*',
   '--allowlist-function', 'glsl_.*',
diff --git a/src/compiler/rust/nir.rs b/src/compiler/rust/nir.rs
index 055948461d9..fae5c1f31ea 100644
--- a/src/compiler/rust/nir.rs
+++ b/src/compiler/rust/nir.rs
@@ -70,6 +70,37 @@ impl<'a, T: 'a> Iterator for ExecListIter<'a, T> {
     }
 }
 
+#[derive(Clone, Copy, Eq, Hash, PartialEq)]
+pub struct ALUType(nir_alu_type);
+
+impl ALUType {
+    pub const INT: Self = Self(nir_type_int);
+    pub const UINT: Self = Self(nir_type_uint);
+    pub const BOOL: Self = Self(nir_type_bool);
+    pub const FLOAT: Self = Self(nir_type_float);
+
+    pub fn new(base: Self, bit_size: u8) -> Self {
+        assert!(bit_size.is_power_of_two());
+        assert!(bit_size & (NIR_ALU_TYPE_BASE_TYPE_MASK as u8) == 0);
+        assert!(
+            base.is_base_type() || bit_size == 0 || bit_size == base.bit_size()
+        );
+        Self(base.0 | bit_size)
+    }
+
+    pub fn is_base_type(&self) -> bool {
+        self.bit_size() == 0
+    }
+
+    pub fn bit_size(&self) -> u8 {
+        self.0 & (NIR_ALU_TYPE_SIZE_MASK as u8)
+    }
+
+    pub fn base_type(&self) -> Self {
+        Self(self.0 & (NIR_ALU_TYPE_BASE_TYPE_MASK as u8))
+    }
+}
+
 impl nir_def {
     pub fn parent_instr<'a>(&'a self) -> &'a nir_instr {
         unsafe { NonNull::new(self.parent_instr).unwrap().as_ref() }