uiua/
function.rs

1use std::{fmt, hash::Hash};
2
3use serde::*;
4
5use crate::{CodeSpan, Ident, Primitive};
6
7/// A function stack signature
8#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default, Serialize, Deserialize)]
9#[serde(from = "(usize, usize)", into = "(usize, usize)")]
10pub struct Signature {
11    /// The number of arguments the function pops off the stack
12    args: u16,
13    /// The number of values the function pushes onto the stack
14    outputs: u16,
15    /// The number of arguments the function pops off the under stack
16    under_args: u16,
17    /// The number of values the function pushes onto the under stack
18    under_outputs: u16,
19}
20
21impl From<(usize, usize)> for Signature {
22    fn from((args, outputs): (usize, usize)) -> Self {
23        Self::new(args, outputs)
24    }
25}
26
27impl From<Signature> for (usize, usize) {
28    fn from(sig: Signature) -> Self {
29        (sig.args(), sig.outputs())
30    }
31}
32
33impl Signature {
34    /// Create a new signature with the given number of arguments and outputs
35    pub const fn new(args: usize, outputs: usize) -> Self {
36        Self {
37            args: args as u16,
38            outputs: outputs as u16,
39            under_args: 0,
40            under_outputs: 0,
41        }
42    }
43    /// Set the number of arguments and outputs of the under stack
44    pub fn with_under(self, under_args: usize, under_outputs: usize) -> Self {
45        Self {
46            args: self.args,
47            outputs: self.outputs,
48            under_args: under_args as u16,
49            under_outputs: under_outputs as u16,
50        }
51    }
52    /// Get the number of arguments
53    #[inline(always)]
54    pub const fn args(&self) -> usize {
55        self.args as usize
56    }
57    /// Get the number of outputs
58    #[inline(always)]
59    pub const fn outputs(&self) -> usize {
60        self.outputs as usize
61    }
62    /// Get the number of under arguments
63    pub const fn under_args(&self) -> usize {
64        self.under_args as usize
65    }
66    /// Get the number of under outputs
67    pub const fn under_outputs(&self) -> usize {
68        self.under_outputs as usize
69    }
70    /// Set the number of arguments
71    pub const fn set_args(&mut self, args: usize) {
72        self.args = args as u16;
73    }
74    /// Set the number of outputs
75    pub const fn set_outputs(&mut self, outputs: usize) {
76        self.outputs = outputs as u16;
77    }
78    /// Update the number of arguments
79    pub fn update_args(&mut self, f: impl FnOnce(usize) -> usize) {
80        self.args = f(self.args()) as u16;
81    }
82    /// Update the number of outputs
83    pub fn update_outputs(&mut self, f: impl FnOnce(usize) -> usize) {
84        self.outputs = f(self.outputs()) as u16;
85    }
86    /// Update the number of under arguments
87    pub fn update_under_args(&mut self, f: impl FnOnce(usize) -> usize) {
88        self.under_args = f(self.under_args()) as u16;
89    }
90    /// Update the number of under outputs
91    pub fn update_under_outputs(&mut self, f: impl FnOnce(usize) -> usize) {
92        self.under_outputs = f(self.under_outputs()) as u16;
93    }
94    /// Update the number of arguments and outputs
95    pub fn update_args_outputs(&mut self, f: impl FnOnce(usize, usize) -> (usize, usize)) {
96        let (args, outputs) = f(self.args(), self.outputs());
97        self.args = args as u16;
98        self.outputs = outputs as u16;
99    }
100    /// Check if this signature changes the stack size by the same amount as another signature
101    pub fn is_compatible_with(self, other: Self) -> bool {
102        self.args as isize - self.outputs as isize == other.args as isize - other.outputs as isize
103    }
104    /// Check if this [`Signature::is_compatible_with`] another signature and has at least as many arguments
105    pub fn is_superset_of(self, other: Self) -> bool {
106        self.is_compatible_with(other) && self.args >= other.args
107    }
108    /// Check if this [`Signature::is_compatible_with`] another signature and has at most as many arguments
109    pub fn is_subset_of(self, other: Self) -> bool {
110        self.is_compatible_with(other) && self.args <= other.args
111    }
112    /// Get the signature that has the maximum of the arguments and outputs of this signature and another
113    pub fn max_with(self, other: Self) -> Self {
114        Self::new(
115            self.args().max(other.args()),
116            self.outputs().max(other.outputs()),
117        )
118        .with_under(
119            self.under_args().max(other.under_args()),
120            self.under_outputs().max(other.under_outputs()),
121        )
122    }
123    /// Compose signatures as if a function with signature `other` was called before a function with signature `self`
124    pub fn compose(self, other: Self) -> Self {
125        let args = other.args() + self.args().saturating_sub(other.outputs());
126        let outputs = self.outputs() + other.outputs().saturating_sub(self.args());
127        let under_args =
128            other.under_args() + self.under_args().saturating_sub(other.under_outputs());
129        let under_outputs =
130            self.under_outputs() + other.under_outputs().saturating_sub(self.under_args());
131        Self::new(args, outputs).with_under(under_args, under_outputs)
132    }
133    /// Get the un-inverse of this signature
134    pub fn inverse(self) -> Self {
135        Self::new(self.outputs(), self.args())
136    }
137    /// The the anti-inverse of this signature
138    pub fn anti(self) -> Option<Self> {
139        if self.args == 0 {
140            return None;
141        }
142        Some(Signature::new(self.outputs() + 1, self.args() - 1))
143    }
144    /// The signature on the under stack
145    pub fn under(self) -> Signature {
146        Signature::new(self.under_args(), self.under_outputs())
147    }
148}
149
150impl PartialEq<(usize, usize)> for Signature {
151    fn eq(&self, other: &(usize, usize)) -> bool {
152        self.args() == other.0 && self.outputs() == other.1
153    }
154}
155
156impl fmt::Debug for Signature {
157    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
158        write!(f, "|{}.{}", self.args, self.outputs)
159    }
160}
161
162impl fmt::Display for Signature {
163    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
164        write!(f, "|{}.{}", self.args, self.outputs)
165    }
166}
167
168/// A function that executes Rust code
169#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
170pub struct DynamicFunction {
171    /// An index used to look up the function
172    pub(crate) index: usize,
173    /// The function's signature
174    pub(crate) sig: Signature,
175}
176
177impl From<(usize, Signature)> for DynamicFunction {
178    fn from((index, sig): (usize, Signature)) -> Self {
179        Self { index, sig }
180    }
181}
182
183impl From<DynamicFunction> for (usize, Signature) {
184    fn from(func: DynamicFunction) -> Self {
185        (func.index, func.sig)
186    }
187}
188
189impl DynamicFunction {
190    /// Get the function's signature
191    pub fn signature(&self) -> Signature {
192        self.sig
193    }
194}
195
196impl fmt::Debug for DynamicFunction {
197    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
198        write!(f, "<dynamic#{:x}>", self.index)
199    }
200}
201
202/// A Uiua function id
203#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
204#[serde(untagged)]
205pub enum FunctionId {
206    /// Just a primitive
207    Primitive(Primitive),
208    /// A named function
209    Named(Ident),
210    /// A macro expansion
211    Macro(Option<Ident>, CodeSpan),
212    /// The top-level function
213    Main,
214    #[doc(hidden)]
215    /// Implementation detail
216    Unnamed,
217}
218
219impl PartialEq<&str> for FunctionId {
220    fn eq(&self, other: &&str) -> bool {
221        match self {
222            FunctionId::Named(name) => &&**name == other,
223            _ => false,
224        }
225    }
226}
227
228impl From<Ident> for FunctionId {
229    fn from(name: Ident) -> Self {
230        Self::Named(name)
231    }
232}
233
234impl From<Primitive> for FunctionId {
235    fn from(op: Primitive) -> Self {
236        Self::Primitive(op)
237    }
238}
239
240impl fmt::Display for FunctionId {
241    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
242        match self {
243            FunctionId::Named(name) => write!(f, "{name}"),
244            FunctionId::Primitive(prim) => write!(f, "{prim}"),
245            FunctionId::Macro(Some(name), span) => write!(f, "macro expansion of {name} at {span}"),
246            FunctionId::Macro(None, span) => write!(f, "macro expansion of at {span}"),
247            FunctionId::Main => write!(f, "main"),
248            FunctionId::Unnamed => write!(f, "unnamed"),
249        }
250    }
251}