leo_ast/stub/
function_stub.rs

1// Copyright (C) 2019-2025 Provable Inc.
2// This file is part of the Leo library.
3
4// The Leo library is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// The Leo library is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
16
17use crate::{
18    Annotation,
19    CompositeType,
20    Function,
21    FutureType,
22    Identifier,
23    Input,
24    Location,
25    Mode,
26    Node,
27    NodeID,
28    Output,
29    Path,
30    ProgramId,
31    TupleType,
32    Type,
33    Variant,
34};
35use leo_span::{Span, Symbol, sym};
36
37use itertools::Itertools;
38use serde::{Deserialize, Serialize};
39use snarkvm::{
40    console::program::{
41        FinalizeType::{Future as FutureFinalizeType, Plaintext as PlaintextFinalizeType},
42        RegisterType::{ExternalRecord, Future, Plaintext, Record},
43    },
44    prelude::{Network, ValueType},
45    synthesizer::program::{ClosureCore, FunctionCore},
46};
47use std::fmt;
48
49/// A function stub definition.
50#[derive(Clone, Serialize, Deserialize)]
51pub struct FunctionStub {
52    /// Annotations on the function.
53    pub annotations: Vec<Annotation>,
54    /// Is this function a transition, inlined, or a regular function?.
55    pub variant: Variant,
56    /// The function identifier, e.g., `foo` in `function foo(...) { ... }`.
57    pub identifier: Identifier,
58    /// The function's input parameters.
59    pub input: Vec<Input>,
60    /// The function's output declarations.
61    pub output: Vec<Output>,
62    /// The function's output type.
63    pub output_type: Type,
64    /// The entire span of the function definition.
65    pub span: Span,
66    /// The ID of the node.
67    pub id: NodeID,
68}
69
70impl PartialEq for FunctionStub {
71    fn eq(&self, other: &Self) -> bool {
72        self.identifier == other.identifier
73    }
74}
75
76impl Eq for FunctionStub {}
77
78impl FunctionStub {
79    /// Initialize a new function.
80    #[allow(clippy::too_many_arguments)]
81    pub fn new(
82        annotations: Vec<Annotation>,
83        _is_async: bool,
84        variant: Variant,
85        identifier: Identifier,
86        input: Vec<Input>,
87        output: Vec<Output>,
88        span: Span,
89        id: NodeID,
90    ) -> Self {
91        let output_type = match output.len() {
92            0 => Type::Unit,
93            1 => output[0].type_.clone(),
94            _ => Type::Tuple(TupleType::new(output.iter().map(|o| o.type_.clone()).collect())),
95        };
96
97        FunctionStub { annotations, variant, identifier, input, output, output_type, span, id }
98    }
99
100    /// Returns function name.
101    pub fn name(&self) -> Symbol {
102        self.identifier.name
103    }
104
105    /// Returns `true` if the function name is `main`.
106    pub fn is_main(&self) -> bool {
107        self.name() == sym::main
108    }
109
110    /// Private formatting method used for optimizing [fmt::Debug] and [fmt::Display] implementations.
111    fn format(&self, f: &mut fmt::Formatter) -> fmt::Result {
112        match self.variant {
113            Variant::Inline => write!(f, "inline ")?,
114            Variant::Script => write!(f, "script ")?,
115            Variant::Function | Variant::AsyncFunction => write!(f, "function ")?,
116            Variant::Transition | Variant::AsyncTransition => write!(f, "transition ")?,
117        }
118        write!(f, "{}", self.identifier)?;
119
120        let parameters = self.input.iter().map(|x| x.to_string()).collect::<Vec<_>>().join(",");
121        let returns = match self.output.len() {
122            0 => "()".to_string(),
123            1 => self.output[0].to_string(),
124            _ => self.output.iter().map(|x| x.to_string()).collect::<Vec<_>>().join(","),
125        };
126        write!(f, "({parameters}) -> {returns}")?;
127
128        Ok(())
129    }
130
131    /// Converts from snarkvm function type to leo FunctionStub, while also carrying the parent program name.
132    pub fn from_function_core<N: Network>(function: &FunctionCore<N>, program: Symbol) -> Self {
133        let outputs = function
134            .outputs()
135            .iter()
136            .map(|output| match output.value_type() {
137                ValueType::Constant(val) => vec![Output {
138                    mode: Mode::Constant,
139                    type_: Type::from_snarkvm(val, None),
140                    span: Default::default(),
141                    id: Default::default(),
142                }],
143                ValueType::Public(val) => vec![Output {
144                    mode: Mode::Public,
145                    type_: Type::from_snarkvm(val, None),
146                    span: Default::default(),
147                    id: Default::default(),
148                }],
149                ValueType::Private(val) => vec![Output {
150                    mode: Mode::Private,
151                    type_: Type::from_snarkvm(val, None),
152                    span: Default::default(),
153                    id: Default::default(),
154                }],
155                ValueType::Record(id) => vec![Output {
156                    mode: Mode::None,
157                    type_: Type::Composite(CompositeType {
158                        path: {
159                            let ident = Identifier::from(id);
160                            Path::from(ident).with_absolute_path(Some(vec![ident.name]))
161                        },
162                        const_arguments: Vec::new(),
163                        program: Some(program),
164                    }),
165                    span: Default::default(),
166                    id: Default::default(),
167                }],
168                ValueType::ExternalRecord(loc) => {
169                    vec![Output {
170                        mode: Mode::None,
171                        span: Default::default(),
172                        id: Default::default(),
173                        type_: Type::Composite(CompositeType {
174                            path: {
175                                let ident = Identifier::from(loc.resource());
176                                Path::from(ident).with_absolute_path(Some(vec![ident.name]))
177                            },
178                            const_arguments: Vec::new(),
179                            program: Some(ProgramId::from(loc.program_id()).name.name),
180                        }),
181                    }]
182                }
183                ValueType::Future(_) => vec![Output {
184                    mode: Mode::None,
185                    span: Default::default(),
186                    id: Default::default(),
187                    type_: Type::Future(FutureType::new(
188                        Vec::new(),
189                        Some(Location::new(program, vec![Symbol::intern(&function.name().to_string())])),
190                        false,
191                    )),
192                }],
193            })
194            .collect_vec()
195            .concat();
196        let output_vec = outputs.iter().map(|output| output.type_.clone()).collect_vec();
197        let output_type = match output_vec.len() {
198            0 => Type::Unit,
199            1 => output_vec[0].clone(),
200            _ => Type::Tuple(TupleType::new(output_vec)),
201        };
202
203        Self {
204            annotations: Vec::new(),
205            variant: match function.finalize_logic().is_some() {
206                true => Variant::AsyncTransition,
207                false => Variant::Transition,
208            },
209            identifier: Identifier::from(function.name()),
210            input: function
211                .inputs()
212                .iter()
213                .enumerate()
214                .map(|(index, input)| {
215                    let arg_name = Identifier::new(Symbol::intern(&format!("arg{}", index + 1)), Default::default());
216                    match input.value_type() {
217                        ValueType::Constant(val) => Input {
218                            identifier: arg_name,
219                            mode: Mode::Constant,
220                            type_: Type::from_snarkvm(val, None),
221                            span: Default::default(),
222                            id: Default::default(),
223                        },
224                        ValueType::Public(val) => Input {
225                            identifier: arg_name,
226                            mode: Mode::Public,
227                            type_: Type::from_snarkvm(val, None),
228                            span: Default::default(),
229                            id: Default::default(),
230                        },
231                        ValueType::Private(val) => Input {
232                            identifier: arg_name,
233                            mode: Mode::Private,
234                            type_: Type::from_snarkvm(val, None),
235                            span: Default::default(),
236                            id: Default::default(),
237                        },
238                        ValueType::Record(id) => Input {
239                            identifier: arg_name,
240                            mode: Mode::None,
241                            type_: Type::Composite(CompositeType {
242                                path: {
243                                    let ident = Identifier::from(id);
244                                    Path::from(ident).with_absolute_path(Some(vec![ident.name]))
245                                },
246                                const_arguments: Vec::new(),
247                                program: Some(program),
248                            }),
249                            span: Default::default(),
250                            id: Default::default(),
251                        },
252                        ValueType::ExternalRecord(loc) => Input {
253                            identifier: arg_name,
254                            mode: Mode::None,
255                            span: Default::default(),
256                            id: Default::default(),
257                            type_: Type::Composite(CompositeType {
258                                path: {
259                                    let ident = Identifier::from(loc.resource());
260                                    Path::from(ident).with_absolute_path(Some(vec![ident.name]))
261                                },
262                                const_arguments: Vec::new(),
263                                program: Some(ProgramId::from(loc.program_id()).name.name),
264                            }),
265                        },
266                        ValueType::Future(_) => panic!("Functions do not contain futures as inputs"),
267                    }
268                })
269                .collect_vec(),
270            output: outputs,
271            output_type,
272            span: Default::default(),
273            id: Default::default(),
274        }
275    }
276
277    pub fn from_finalize<N: Network>(function: &FunctionCore<N>, key_name: Symbol, program: Symbol) -> Self {
278        Self {
279            annotations: Vec::new(),
280            variant: Variant::AsyncFunction,
281            identifier: Identifier::new(key_name, Default::default()),
282            input: function
283                .finalize_logic()
284                .unwrap()
285                .inputs()
286                .iter()
287                .enumerate()
288                .map(|(index, input)| Input {
289                    identifier: Identifier::new(Symbol::intern(&format!("arg{}", index + 1)), Default::default()),
290                    mode: Mode::None,
291                    type_: match input.finalize_type() {
292                        PlaintextFinalizeType(val) => Type::from_snarkvm(val, Some(program)),
293                        FutureFinalizeType(val) => Type::Future(FutureType::new(
294                            Vec::new(),
295                            Some(Location::new(Identifier::from(val.program_id().name()).name, vec![Symbol::intern(
296                                &format!("finalize/{}", val.resource()),
297                            )])),
298                            false,
299                        )),
300                    },
301                    span: Default::default(),
302                    id: Default::default(),
303                })
304                .collect_vec(),
305            output: Vec::new(),
306            output_type: Type::Unit,
307            span: Default::default(),
308            id: 0,
309        }
310    }
311
312    pub fn from_closure<N: Network>(closure: &ClosureCore<N>, program: Symbol) -> Self {
313        let outputs = closure
314            .outputs()
315            .iter()
316            .map(|output| match output.register_type() {
317                Plaintext(val) => Output {
318                    mode: Mode::None,
319                    type_: Type::from_snarkvm(val, Some(program)),
320                    span: Default::default(),
321                    id: Default::default(),
322                },
323                Record(_) => panic!("Closures do not return records"),
324                ExternalRecord(_) => panic!("Closures do not return external records"),
325                Future(_) => panic!("Closures do not return futures"),
326            })
327            .collect_vec();
328        let output_vec = outputs.iter().map(|output| output.type_.clone()).collect_vec();
329        let output_type = match output_vec.len() {
330            0 => Type::Unit,
331            1 => output_vec[0].clone(),
332            _ => Type::Tuple(TupleType::new(output_vec)),
333        };
334        Self {
335            annotations: Vec::new(),
336            variant: Variant::Function,
337            identifier: Identifier::from(closure.name()),
338            input: closure
339                .inputs()
340                .iter()
341                .enumerate()
342                .map(|(index, input)| {
343                    let arg_name = Identifier::new(Symbol::intern(&format!("arg{}", index + 1)), Default::default());
344                    match input.register_type() {
345                        Plaintext(val) => Input {
346                            identifier: arg_name,
347                            mode: Mode::None,
348                            type_: Type::from_snarkvm(val, None),
349                            span: Default::default(),
350                            id: Default::default(),
351                        },
352                        Record(_) => panic!("Closures do not contain records as inputs"),
353                        ExternalRecord(_) => panic!("Closures do not contain external records as inputs"),
354                        Future(_) => panic!("Closures do not contain futures as inputs"),
355                    }
356                })
357                .collect_vec(),
358            output: outputs,
359            output_type,
360            span: Default::default(),
361            id: Default::default(),
362        }
363    }
364}
365
366impl From<Function> for FunctionStub {
367    fn from(function: Function) -> Self {
368        Self {
369            annotations: function.annotations,
370            variant: function.variant,
371            identifier: function.identifier,
372            input: function.input,
373            output: function.output,
374            output_type: function.output_type,
375            span: function.span,
376            id: function.id,
377        }
378    }
379}
380
381impl fmt::Debug for FunctionStub {
382    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
383        self.format(f)
384    }
385}
386
387impl fmt::Display for FunctionStub {
388    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
389        self.format(f)
390    }
391}
392
393crate::simple_node_impl!(FunctionStub);
leo_ast/stub/function_stub.rs

leo_ast/stub/
function_stub.rs
