Type inference for external finalize inputs

compiler/ast/src/stub/function_stub.rs

@@ -35,7 +35,7 @@ use crate::{
 };
 use leo_span::{sym, Span, Symbol};
 
-use crate::Type::Composite;
+use crate::Type::{Composite};
 use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 use snarkvm::{
@@ -47,6 +47,7 @@ use snarkvm::{
     synthesizer::program::{ClosureCore, CommandTrait, FunctionCore, InstructionTrait},
 };
 use std::fmt;
+use crate::stub::future_stub::FutureStub;
 
 /// A function stub definition.
 #[derive(Clone, Serialize, Deserialize)]
@@ -59,6 +60,8 @@ pub struct FunctionStub {
     pub variant: Variant,
     /// The function identifier, e.g., `foo` in `function foo(...) { ... }`.
     pub identifier: Identifier,
+    /// Ordered list of futures inputted to finalize.
+    pub future_stubs: Vec<FutureStub>,
     /// The function's input parameters.
     pub input: Vec<Input>,
     /// The function's output declarations.
@@ -104,7 +107,7 @@ impl FunctionStub {
             _ => Type::Tuple(TupleType::new(output.iter().map(get_output_type).collect())),
         };
 
-        FunctionStub { annotations, is_async, variant, identifier, input, output, output_type, span, id }
+        FunctionStub { annotations, is_async, variant, identifier, future_stubs: Vec::new(), input, output, output_type, span, id }
     }
 
     /// Returns function name.
@@ -205,6 +208,7 @@ impl FunctionStub {
             is_async: function.finalize_logic().is_some(),
             variant: Variant::Transition,
             identifier: Identifier::from(function.name()),
+            future_stubs: Vec::new(),
             input: function
                 .inputs()
                 .iter()
@@ -267,6 +271,10 @@ impl FunctionStub {
             is_async: true,
             variant: Variant::Transition,
             identifier: Identifier::new(name, Default::default()),
+            future_stubs: function.inputs().iter().filter_map(|input| match input.value_type() {
+                ValueType::Future(val) => Some(FutureStub::new(Identifier::from(val.program_id().name()).name, Identifier::from(val.resource()).name)),
+                _ => None,
+            }).collect(),
             input: function
                 .finalize_logic()
                 .unwrap()
@@ -288,11 +296,11 @@ impl FunctionStub {
                 .collect_vec(),
             output: vec![Output::Internal(FunctionOutput {
                 mode: Mode::Public,
-                type_: Type::Future(FutureType { inputs: None }),
+                type_: Type::Future(FutureType { inputs: Vec::new() }),
                 span: Default::default(),
                 id: 0,
             })],
-            output_type: Type::Future(FutureType { inputs: None }),
+            output_type: Type::Future(FutureType { inputs: Vec::new() }),
             span: Default::default(),
             id: 0,
         }
@@ -334,6 +342,7 @@ impl FunctionStub {
             is_async: false,
             variant: Variant::Standard,
             identifier: Identifier::from(closure.name()),
+            future_stubs: Vec::new(),
             input: closure
                 .inputs()
                 .iter()
@@ -369,6 +378,7 @@ impl From<Function> for FunctionStub {
             is_async: function.is_async,
             variant: function.variant,
             identifier: function.identifier,
+            future_stubs: Vec::new(),
             input: function.input,
             output: function.output,
             output_type: function.output_type,

compiler/ast/src/stub/future_stub.rs

@@ -0,0 +1,61 @@
+// Copyright (C) 2019-2023 Aleo Systems Inc.
+// This file is part of the Leo library.
+
+// The Leo library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// The Leo library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
+
+use leo_span::Symbol;
+
+use serde::{Deserialize, Serialize};
+
+/// A future stub definition.
+#[derive(Clone, Serialize, Deserialize)]
+pub struct FutureStub {
+    program: Symbol,
+    function: Symbol,
+}
+
+impl PartialEq for FutureStub {
+    fn eq(&self, other: &Self) -> bool {
+        self.program == other.program && self.function == other.function
+    }
+}
+
+impl Eq for FutureStub {}
+
+impl FutureStub {
+    /// Initialize a new future stub.
+    pub fn new(
+        program: Symbol,
+        function: Symbol,
+    ) -> Self {
+        FutureStub {
+            program,
+            function,
+        }
+    }
+    
+    pub fn to_key(&self) -> (Symbol, Symbol) {
+        (self.program, self.function)
+    }
+    
+    /// Get the program.
+    pub fn program(&self) -> Symbol {
+        self.program
+    }
+    
+    /// Get the function.
+    pub fn function(&self) -> Symbol {
+        self.function
+    }
+}

compiler/ast/src/stub/mod.rs

@@ -19,6 +19,9 @@
 pub mod function_stub;
 pub use function_stub::*;
 
+pub mod future_stub;
+pub use future_stub;
+
 use crate::{Composite, ConstDeclaration, Identifier, Mapping, NodeID, ProgramId};
 use leo_span::{Span, Symbol};
 use serde::{Deserialize, Serialize};

compiler/ast/src/types/future.rs

@@ -27,6 +27,10 @@ pub struct FutureType {
 }
 
 impl FutureType {
+    /// Initialize a new future type.
+    pub fn new(inputs: Vec<Type>) -> Self {
+        Self { inputs }
+    }
     /// Returns the inputs of the future type.
     pub fn inputs(&self) -> &[Type] {
         &self.inputs

compiler/passes/src/type_checking/check_program.rs

@@ -25,6 +25,7 @@ use snarkvm::console::network::{Network, Testnet3};
 use indexmap::IndexSet;
 use leo_ast::Input::{External, Internal};
 use std::collections::HashSet;
+use leo_ast::Type::Future;
 
 // TODO: Cleanup logic for tuples.
 
@@ -32,6 +33,7 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
     fn visit_program(&mut self, input: &'a Program) {
         // Typecheck the program's stubs.
         input.stubs.iter().for_each(|(symbol, stub)| {
+            // Check that naming and ordering is consistent.
             if symbol != &stub.stub_id.name.name {
                 self.emit_err(TypeCheckerError::stub_name_mismatch(
                     symbol,
@@ -84,6 +86,25 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
         // Create a new child scope for the function's parameters and body.
         let scope_index = self.create_child_scope();
 
+        // Create future stubs.
+        let finalize_input_map = &mut self.finalize_input_types;
+        let mut future_stubs = input.future_stubs.clone();
+        let resolved_inputs = input.input.iter().map(|input_mode| {
+                match input_mode {
+                    Internal(function_input) => match &function_input.type_ {
+                        Future(_) => {
+                            // Since we traverse stubs in post-order, we can assume that the corresponding finalize stub has already been traversed.
+                            Future(FutureType::new(finalize_input_map.get(&future_stubs.pop().unwrap().to_key()).unwrap().clone()))
+                        }
+                        _ => function_input.clone().type_,
+                    },
+                    External(_) => {}
+                }
+            }).collect();
+        assert!(future_stubs.is_empty(), "Disassembler produced malformed stub.");
+
+        finalize_input_map.insert((self.scope_state.program_name.unwrap(), input.identifier.name), resolved_inputs);
+
         // Query helper function to type check function parameters and outputs.
         self.check_function_signature(&Function::from(input.clone()));
 
@@ -339,7 +360,7 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
             self.emit_err(TypeCheckerError::async_transition_must_call_async_function(function.span));
         }
 
-        // Check that all futures were awaited exactly once. 
+        // Check that all futures were awaited exactly once.
         if self.scope_state.is_finalize {
             // Throw error if not all futures awaits even appear once.
             if !self.await_checker.static_to_await.is_empty() {

compiler/passes/src/type_checking/checker.rs

@@ -16,21 +16,7 @@
 
 use crate::{CallGraph, StructGraph, SymbolTable, TreeNode, TypeTable, VariableSymbol, VariableType};
 
-use leo_ast::{
-    Composite,
-    CompositeType,
-    CoreConstant,
-    CoreFunction,
-    Function,
-    Identifier,
-    IntegerType,
-    MappingType,
-    Mode,
-    Node,
-    Output,
-    Type,
-    Variant,
-};
+use leo_ast::{Composite, CompositeType, CoreConstant, CoreFunction, Function, Identifier, Input, IntegerType, MappingType, Mode, Node, Output, Type, Variant};
 use leo_errors::{emitter::Handler, TypeCheckerError, TypeCheckerWarning};
 use leo_span::{Span, Symbol};
 
@@ -57,7 +43,7 @@ pub struct TypeChecker<'a> {
     /// Struct to store the state relevant to checking all futures are awaited.
     pub(crate) await_checker: AwaitChecker,
     /// Mapping from async function name to the inferred input types.
-    pub(crate) finalize_input_types: IndexMap<Symbol, Vec<Type>>,
+    pub(crate) finalize_input_types: IndexMap<(Symbol, Symbol), Vec<Type>>,
 }
 
 const ADDRESS_TYPE: Type = Type::Address;
@@ -1243,8 +1229,8 @@ impl<'a> TypeChecker<'a> {
     pub(crate) fn check_function_signature(&mut self, function: &Function) {
         self.scope_state.variant = Some(function.variant);
 
-        // Special type checking for finalize blocks.
-        if self.scope_state.is_finalize {
+        // Special type checking for finalize blocks. Can skip for stubs.
+        if self.scope_state.is_finalize & !self.scope_state.is_stub {
             if let Some(inferred_future_types) =
                 self.finalize_input_types.borrow().get(&self.scope_state.function.unwrap())
             {

utils/disassembler/src/lib.rs

@@ -75,7 +75,7 @@ pub fn disassemble<N: Network, Instruction: InstructionTrait<N>, Command: Comman
                             "finalize/{}",
                             Symbol::intern(&Identifier::from(id).name.to_string())
                         ));
-                        Some((name, FunctionStub::from_finalize(f, name)))
+                        Some((name, FunctionStub::from_finalize(function, name)))
                     }
                     None => None,
                 })