leo_ast/stub/
function_stub.rs

1// Copyright (C) 2019-2026 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::RegisterType,
41    prelude::{FinalizeType, Network, ValueType},
42    synthesizer::program::{ClosureCore, FunctionCore},
43};
44use std::fmt;
45
46/// A function stub definition.
47#[derive(Clone, Serialize, Deserialize)]
48pub struct FunctionStub {
49    /// Annotations on the function.
50    pub annotations: Vec<Annotation>,
51    /// Is this function a transition, inlined, or a regular function?.
52    pub variant: Variant,
53    /// The function identifier, e.g., `foo` in `function foo(...) { ... }`.
54    pub identifier: Identifier,
55    /// The function's input parameters.
56    pub input: Vec<Input>,
57    /// The function's output declarations.
58    pub output: Vec<Output>,
59    /// The function's output type.
60    pub output_type: Type,
61    /// The entire span of the function definition.
62    pub span: Span,
63    /// The ID of the node.
64    pub id: NodeID,
65}
66
67impl PartialEq for FunctionStub {
68    fn eq(&self, other: &Self) -> bool {
69        self.identifier == other.identifier
70    }
71}
72
73impl Eq for FunctionStub {}
74
75impl FunctionStub {
76    /// Initialize a new function.
77    #[allow(clippy::too_many_arguments)]
78    pub fn new(
79        annotations: Vec<Annotation>,
80        _is_async: bool,
81        variant: Variant,
82        identifier: Identifier,
83        input: Vec<Input>,
84        output: Vec<Output>,
85        span: Span,
86        id: NodeID,
87    ) -> Self {
88        let output_type = match output.len() {
89            0 => Type::Unit,
90            1 => output[0].type_.clone(),
91            _ => Type::Tuple(TupleType::new(output.iter().map(|o| o.type_.clone()).collect())),
92        };
93
94        FunctionStub { annotations, variant, identifier, input, output, output_type, span, id }
95    }
96
97    /// Returns function name.
98    pub fn name(&self) -> Symbol {
99        self.identifier.name
100    }
101
102    /// Returns `true` if the function name is `main`.
103    pub fn is_main(&self) -> bool {
104        self.name() == sym::main
105    }
106
107    /// Returns `true` if any output of the function is a `Final`
108    pub fn has_final_output(&self) -> bool {
109        self.output.iter().any(|o| matches!(o.type_, Type::Future(_)))
110    }
111
112    /// Private formatting method used for optimizing [fmt::Debug] and [fmt::Display] implementations.
113    fn format(&self, f: &mut fmt::Formatter) -> fmt::Result {
114        match self.variant {
115            Variant::FinalFn => write!(f, "final fn ")?,
116            Variant::Finalize => write!(f, "finalize ")?,
117            Variant::Fn => write!(f, "fn ")?,
118            Variant::EntryPoint => write!(f, "entry ")?,
119        }
120        write!(f, "{}", self.identifier)?;
121
122        let parameters = self.input.iter().map(|x| x.to_string()).collect::<Vec<_>>().join(",");
123        let returns = match self.output.len() {
124            0 => "()".to_string(),
125            1 => self.output[0].to_string(),
126            _ => self.output.iter().map(|x| x.to_string()).collect::<Vec<_>>().join(","),
127        };
128        write!(f, "({parameters}) -> {returns}")?;
129
130        Ok(())
131    }
132
133    /// Converts from snarkvm function type to leo FunctionStub, while also carrying the parent program name.
134    pub fn from_function_core<N: Network>(function: &FunctionCore<N>, program_id: ProgramId) -> Self {
135        let outputs = function
136            .outputs()
137            .iter()
138            .map(|output| match output.value_type() {
139                ValueType::Constant(val) => vec![Output {
140                    mode: Mode::Constant,
141                    type_: Type::from_snarkvm(val, program_id),
142                    span: Default::default(),
143                    id: Default::default(),
144                }],
145                ValueType::Public(val) => vec![Output {
146                    mode: Mode::Public,
147                    type_: Type::from_snarkvm(val, program_id),
148                    span: Default::default(),
149                    id: Default::default(),
150                }],
151                ValueType::Private(val) => vec![Output {
152                    mode: Mode::Private,
153                    type_: Type::from_snarkvm(val, program_id),
154                    span: Default::default(),
155                    id: Default::default(),
156                }],
157                ValueType::Record(id) => vec![Output {
158                    mode: Mode::None,
159                    type_: Type::Composite(CompositeType {
160                        path: {
161                            let ident = Identifier::from(id);
162                            Path::from(ident)
163                                .to_global(Location::new(program_id.as_symbol(), vec![ident.name]))
164                                .with_user_program(program_id)
165                        },
166                        const_arguments: Vec::new(),
167                    }),
168                    span: Default::default(),
169                    id: Default::default(),
170                }],
171                ValueType::ExternalRecord(loc) => {
172                    let external_program_id = ProgramId::from(loc.program_id());
173                    vec![Output {
174                        mode: Mode::None,
175                        span: Default::default(),
176                        id: Default::default(),
177                        type_: Type::Composite(CompositeType {
178                            path: {
179                                let ident = Identifier::from(loc.resource());
180                                Path::from(ident)
181                                    .to_global(Location::new(external_program_id.as_symbol(), vec![ident.name]))
182                                    .with_user_program(external_program_id)
183                            },
184                            const_arguments: Vec::new(),
185                        }),
186                    }]
187                }
188                ValueType::Future(_) => vec![Output {
189                    mode: Mode::None,
190                    span: Default::default(),
191                    id: Default::default(),
192                    type_: Type::Future(FutureType::new(
193                        Vec::new(),
194                        Some(Location::new(program_id.as_symbol(), vec![Symbol::intern(&function.name().to_string())])),
195                        false,
196                    )),
197                }],
198                ValueType::DynamicRecord => vec![Output {
199                    mode: Mode::None,
200                    span: Default::default(),
201                    id: Default::default(),
202                    type_: Type::DynRecord,
203                }],
204                ValueType::DynamicFuture => vec![Output {
205                    mode: Mode::None,
206                    span: Default::default(),
207                    id: Default::default(),
208                    type_: Type::Future(FutureType::new(
209                        Vec::new(),
210                        Some(Location::new(program_id.as_symbol(), vec![Symbol::intern(&function.name().to_string())])),
211                        false,
212                    )),
213                }],
214            })
215            .collect_vec()
216            .concat();
217        let output_vec = outputs.iter().map(|output| output.type_.clone()).collect_vec();
218        let output_type = match output_vec.len() {
219            0 => Type::Unit,
220            1 => output_vec[0].clone(),
221            _ => Type::Tuple(TupleType::new(output_vec)),
222        };
223
224        Self {
225            annotations: Vec::new(),
226            variant: Variant::EntryPoint,
227            identifier: Identifier::from(function.name()),
228            input: function
229                .inputs()
230                .iter()
231                .enumerate()
232                .map(|(index, input)| {
233                    let arg_name = Identifier::new(Symbol::intern(&format!("arg{}", index + 1)), Default::default());
234                    match input.value_type() {
235                        ValueType::Constant(val) => Input {
236                            identifier: arg_name,
237                            mode: Mode::Constant,
238                            type_: Type::from_snarkvm(val, program_id),
239                            span: Default::default(),
240                            id: Default::default(),
241                        },
242                        ValueType::Public(val) => Input {
243                            identifier: arg_name,
244                            mode: Mode::Public,
245                            type_: Type::from_snarkvm(val, program_id),
246                            span: Default::default(),
247                            id: Default::default(),
248                        },
249                        ValueType::Private(val) => Input {
250                            identifier: arg_name,
251                            mode: Mode::Private,
252                            type_: Type::from_snarkvm(val, program_id),
253                            span: Default::default(),
254                            id: Default::default(),
255                        },
256                        ValueType::Record(id) => Input {
257                            identifier: arg_name,
258                            mode: Mode::None,
259                            type_: Type::Composite(CompositeType {
260                                path: {
261                                    let ident = Identifier::from(id);
262                                    Path::from(ident)
263                                        .to_global(Location::new(program_id.as_symbol(), vec![ident.name]))
264                                        .with_user_program(program_id)
265                                },
266                                const_arguments: Vec::new(),
267                            }),
268                            span: Default::default(),
269                            id: Default::default(),
270                        },
271                        ValueType::ExternalRecord(loc) => {
272                            let external_program = ProgramId::from(loc.program_id());
273                            Input {
274                                identifier: arg_name,
275                                mode: Mode::None,
276                                span: Default::default(),
277                                id: Default::default(),
278                                type_: Type::Composite(CompositeType {
279                                    path: {
280                                        let ident = Identifier::from(loc.resource());
281                                        Path::from(ident)
282                                            .to_global(Location::new(external_program.as_symbol(), vec![ident.name]))
283                                            .with_user_program(external_program)
284                                    },
285                                    const_arguments: Vec::new(),
286                                }),
287                            }
288                        }
289                        ValueType::Future(_) | ValueType::DynamicFuture => {
290                            panic!("Functions do not contain futures as inputs")
291                        }
292
293                        ValueType::DynamicRecord => Input {
294                            identifier: arg_name,
295                            mode: Mode::None,
296                            span: Default::default(),
297                            id: Default::default(),
298                            type_: Type::DynRecord,
299                        },
300                    }
301                })
302                .collect_vec(),
303            output: outputs,
304            output_type,
305            span: Default::default(),
306            id: Default::default(),
307        }
308    }
309
310    pub fn from_finalize<N: Network>(function: &FunctionCore<N>, key_name: Symbol, program_id: ProgramId) -> Self {
311        Self {
312            annotations: Vec::new(),
313            variant: Variant::Finalize,
314            identifier: Identifier::new(key_name, Default::default()),
315            input: function
316                .finalize_logic()
317                .unwrap()
318                .inputs()
319                .iter()
320                .enumerate()
321                .map(|(index, input)| Input {
322                    identifier: Identifier::new(Symbol::intern(&format!("arg{}", index + 1)), Default::default()),
323                    mode: Mode::None,
324                    type_: match input.finalize_type() {
325                        FinalizeType::Plaintext(val) => Type::from_snarkvm(val, program_id),
326                        FinalizeType::Future(val) => Type::Future(FutureType::new(
327                            Vec::new(),
328                            Some(Location::new(ProgramId::from(val.program_id()).as_symbol(), vec![Symbol::intern(
329                                &format!("finalize/{}", val.resource()),
330                            )])),
331                            false,
332                        )),
333                        FinalizeType::DynamicFuture => Type::Future(FutureType::new(Vec::new(), None, false)),
334                    },
335                    span: Default::default(),
336                    id: Default::default(),
337                })
338                .collect_vec(),
339            output: Vec::new(),
340            output_type: Type::Unit,
341            span: Default::default(),
342            id: 0,
343        }
344    }
345
346    pub fn from_closure<N: Network>(closure: &ClosureCore<N>, program_id: ProgramId) -> Self {
347        let outputs = closure
348            .outputs()
349            .iter()
350            .map(|output| match output.register_type() {
351                RegisterType::Plaintext(val) => Output {
352                    mode: Mode::None,
353                    type_: Type::from_snarkvm(val, program_id),
354                    span: Default::default(),
355                    id: Default::default(),
356                },
357                RegisterType::Record(_) | RegisterType::DynamicRecord => panic!("Closures do not return records"),
358                RegisterType::ExternalRecord(_) => panic!("Closures do not return external records"),
359                RegisterType::Future(_) | RegisterType::DynamicFuture => panic!("Closures do not return futures"),
360            })
361            .collect_vec();
362        let output_vec = outputs.iter().map(|output| output.type_.clone()).collect_vec();
363        let output_type = match output_vec.len() {
364            0 => Type::Unit,
365            1 => output_vec[0].clone(),
366            _ => Type::Tuple(TupleType::new(output_vec)),
367        };
368        Self {
369            annotations: Vec::new(),
370            variant: Variant::Fn,
371            identifier: Identifier::from(closure.name()),
372            input: closure
373                .inputs()
374                .iter()
375                .enumerate()
376                .map(|(index, input)| {
377                    let arg_name = Identifier::new(Symbol::intern(&format!("arg{}", index + 1)), Default::default());
378                    match input.register_type() {
379                        RegisterType::Plaintext(val) => Input {
380                            identifier: arg_name,
381                            mode: Mode::None,
382                            type_: Type::from_snarkvm(val, program_id),
383                            span: Default::default(),
384                            id: Default::default(),
385                        },
386                        RegisterType::Record(_) | RegisterType::DynamicRecord => {
387                            panic!("Closures do not contain records as inputs")
388                        }
389                        RegisterType::ExternalRecord(_) => panic!("Closures do not contain external records as inputs"),
390                        RegisterType::Future(_) | RegisterType::DynamicFuture => {
391                            panic!("Closures do not contain futures as inputs")
392                        }
393                    }
394                })
395                .collect_vec(),
396            output: outputs,
397            output_type,
398            span: Default::default(),
399            id: Default::default(),
400        }
401    }
402}
403
404impl From<Function> for FunctionStub {
405    fn from(function: Function) -> Self {
406        Self {
407            annotations: function.annotations,
408            variant: function.variant,
409            identifier: function.identifier,
410            input: function.input,
411            output: function.output,
412            output_type: function.output_type,
413            span: function.span,
414            id: function.id,
415        }
416    }
417}
418
419impl fmt::Debug for FunctionStub {
420    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
421        self.format(f)
422    }
423}
424
425impl fmt::Display for FunctionStub {
426    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
427        self.format(f)
428    }
429}
430
431crate::simple_node_impl!(FunctionStub);
leo_ast/stub/function_stub.rs

leo_ast/stub/
function_stub.rs
