Skip to main content

enki_anu/nam_args_api/
dispatch_map.rs

1use std::ops::Range;
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub enum ExpectedParamKind {
5    PerCell,
6    Slice,
7    Atomic,
8    TileScratchpad,
9    ByValue,
10}
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub struct NamParamMeta {
14    pub name: &'static str,
15    pub type_str: &'static str,
16    pub line: u32,
17    pub column: u32,
18    pub is_mutable: bool,
19    pub expected_kind: ExpectedParamKind,
20}
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub struct NamSignatureContract {
24    pub nam_name: &'static str,
25    pub file_path: &'static str,
26    pub line: u32,
27    pub params: &'static [NamParamMeta],
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
31pub enum DispatchSafetyMode {
32    #[default]
33    Safe,
34    Unchecked,
35}
36
37impl DispatchSafetyMode {
38    #[inline(always)]
39    pub fn is_safe(&self) -> bool {
40        matches!(self, Self::Safe)
41    }
42
43    #[inline(always)]
44    pub fn is_unchecked(&self) -> bool {
45        matches!(self, Self::Unchecked)
46    }
47}
48
49#[derive(Debug, Clone, PartialEq, Eq)]
50pub enum ResourceAccessKind {
51    PerCell {
52        element_count: usize,
53    },
54    Slice {
55        element_range: Range<usize>,
56        total_container_len: usize,
57    },
58    Atomic {
59        element_count: usize,
60    },
61    TileScratchpad {
62        element_count: usize,
63    },
64    ValueUniform {
65        byte_size: usize,
66    },
67}
68
69#[derive(Debug, Clone, PartialEq, Eq)]
70pub struct InputResourceRecord {
71    pub arg_index: usize,
72    pub root_id: Option<usize>,
73    pub access_kind: ResourceAccessKind,
74    pub is_mutable: bool,
75    pub element_type_name: &'static str,
76}
77
78impl InputResourceRecord {
79    #[inline]
80    pub fn as_slice_range(&self) -> Option<&Range<usize>> {
81        match &self.access_kind {
82            ResourceAccessKind::Slice { element_range, .. } => Some(element_range),
83            _ => None,
84        }
85    }
86
87    #[inline(always)]
88    pub fn is_write(&self) -> bool {
89        self.is_mutable
90    }
91}
92
93#[derive(Debug, Clone, Copy, PartialEq, Eq)]
94pub struct SpaceContract {
95    pub dimensions: (usize, usize, usize),
96    pub total_cells: usize,
97}
98
99impl SpaceContract {
100    pub fn new(x: usize, y: usize, z: usize) -> Self {
101        let x = x.max(1);
102        let y = y.max(1);
103        let z = z.max(1);
104        let total_cells = x * y * z;
105
106        Self {
107            dimensions: (x, y, z),
108            total_cells,
109        }
110    }
111}
112
113#[derive(Debug, Clone)]
114pub struct NamDispatchMap {
115    pub nam_name: String,
116
117    pub safety_mode: DispatchSafetyMode,
118
119    pub space: SpaceContract,
120
121    pub inputs: Vec<InputResourceRecord>,
122
123    pub call_site: Option<(&'static str, u32, u32)>,
124
125    pub expected_contract: Option<&'static NamSignatureContract>,
126}
127
128impl NamDispatchMap {
129    pub fn new(
130        nam_name: String,
131        safety_mode: DispatchSafetyMode,
132        space: SpaceContract,
133        call_site: Option<(&'static str, u32, u32)>,
134    ) -> Self {
135        Self {
136            nam_name,
137            safety_mode,
138            space,
139            inputs: Vec::with_capacity(8),
140            call_site,
141            expected_contract: None,
142        }
143    }
144
145    #[inline]
146    pub fn with_contract(mut self, contract: &'static NamSignatureContract) -> Self {
147        self.expected_contract = Some(contract);
148        self
149    }
150
151    pub fn push_input(&mut self, input: InputResourceRecord) {
152        self.inputs.push(input);
153    }
154
155    pub fn inputs_for_root(
156        &self,
157        target_root: usize,
158    ) -> impl Iterator<Item = &InputResourceRecord> {
159        self.inputs
160            .iter()
161            .filter(move |i| i.root_id == Some(target_root))
162    }
163
164    pub fn has_mutable_inputs(&self) -> bool {
165        self.inputs.iter().any(|i| i.is_mutable)
166    }
167}