simplify init and post upgrade mod and build files

demergent-labs · Apr 1, 2024 · f7fec1d · f7fec1d
1 parent 5f3e170
commit f7fec1d
Show file tree

Hide file tree

Showing 7 changed files with 81 additions and 122 deletions.
diff --git a/examples/pre_and_post_upgrade/src/main.did b/examples/pre_and_post_upgrade/src/main.did
@@ -2,4 +2,4 @@ type Entry = record { key : text; value : nat64 };
 service : () -> {
   get_entries : () -> (vec Entry) query;
   set_entry : (Entry) -> ();
-}
+}
diff --git a/kybra/compiler/kybra_generate/src/canister_method/init_method/build.rs b/kybra/compiler/kybra_generate/src/canister_method/init_method/build.rs
@@ -1,3 +1,5 @@
+// TODO init_method/build.rs and post_upgrade_method/build.rs are almost identical
+
 use super::rust;
 use cdk_framework::{
     act::node::canister_method::{CanisterMethodType, InitMethod},
@@ -39,6 +41,16 @@ impl PyAst {
     }
 }
 
+fn build_init_methods(
+    init_function_defs: &Vec<SourceMapped<&Located<StmtKind>>>,
+    entry_module_name: &str,
+) -> Result<Vec<InitMethod>, Vec<Error>> {
+    init_function_defs
+        .iter()
+        .map(|init_function_def| build_init_method(Some(init_function_def), entry_module_name))
+        .collect_results()
+}
+
 fn build_init_method(
     init_function_def_option: Option<&SourceMapped<&Located<StmtKind>>>,
     entry_module_name: &str,
@@ -70,13 +82,3 @@ fn build_init_method(
         }),
     }
 }
-
-fn build_init_methods(
-    init_function_defs: &Vec<SourceMapped<&Located<StmtKind>>>,
-    entry_module_name: &str,
-) -> Result<Vec<InitMethod>, Vec<Error>> {
-    init_function_defs
-        .iter()
-        .map(|init_function_def| build_init_method(Some(init_function_def), entry_module_name))
-        .collect_results()
-}
diff --git a/kybra/compiler/kybra_generate/src/canister_method/init_method/mod.rs b/kybra/compiler/kybra_generate/src/canister_method/init_method/mod.rs
@@ -1,2 +1,2 @@
 pub mod build;
-mod rust;
+pub mod rust;
diff --git a/kybra/compiler/kybra_generate/src/canister_method/init_method/rust.rs b/kybra/compiler/kybra_generate/src/canister_method/init_method/rust.rs
@@ -2,7 +2,6 @@ use proc_macro2::TokenStream;
 use quote::quote;
 use rustpython_parser::ast::{Located, StmtKind};
 
-use crate::canister_method::post_upgrade_method::generate_call;
 use crate::canister_method::post_upgrade_method::rust::{
     generate_code_init, generate_ic_object_init, generate_interpreter_init, generate_randomness,
     generate_save_global_interpreter,
@@ -47,3 +46,12 @@ pub fn generate(
         #randomness
     })
 }
+
+pub fn generate_call(
+    function_def_option: &Option<&SourceMapped<&Located<StmtKind>>>,
+) -> Result<TokenStream, Vec<Error>> {
+    match &function_def_option {
+        Some(function_def) => function_def.generate_call_to_py_function(),
+        None => Ok(quote!()),
+    }
+}
diff --git a/kybra/compiler/kybra_generate/src/canister_method/post_upgrade_method/build.rs b/kybra/compiler/kybra_generate/src/canister_method/post_upgrade_method/build.rs
@@ -1,19 +1,16 @@
+use super::rust;
 use cdk_framework::{
-    act::node::{
-        canister_method::{CanisterMethodType, PostUpgradeMethod},
-        Param,
-    },
+    act::node::canister_method::{CanisterMethodType, PostUpgradeMethod},
     traits::CollectResults,
 };
-use proc_macro2::TokenStream;
 use rustpython_parser::ast::{Located, StmtKind};
 
-use super::rust;
 use crate::{
     canister_method::{
         self,
         errors::{MultipleSystemMethods, ReturnTypeMustBeVoid},
     },
+    errors::CollectResults as OtherCollectResults,
     method_utils::params::InternalOrExternal,
     py_ast::PyAst,
     source_map::SourceMapped,
@@ -22,103 +19,71 @@ use crate::{
 
 impl PyAst {
     pub fn build_post_upgrade_method(&self) -> Result<PostUpgradeMethod, Vec<Error>> {
-        let init_function_defs = self.get_canister_stmt_of_type(CanisterMethodType::Init);
         let post_upgrade_function_defs =
             self.get_canister_stmt_of_type(CanisterMethodType::PostUpgrade);
-        let (init_params, post_upgrade_params, post_upgrade_method_body) = (
-            build_init_params(&init_function_defs),
-            build_post_upgrade_info(&post_upgrade_function_defs),
-            generate_post_upgrade_method_body(
-                &self.entry_module_name,
-                post_upgrade_function_defs.get(0),
-            ),
-        )
-            .collect_results()?;
-        let post_upgrade_method_params =
-            select_post_upgrade_method_params(init_params, post_upgrade_params);
 
-        Ok(PostUpgradeMethod {
-            params: post_upgrade_method_params,
-            body: post_upgrade_method_body,
-        })
-    }
-}
+        let post_upgrade_methods =
+            build_post_upgrade_methods(&post_upgrade_function_defs, &self.entry_module_name);
 
-fn build_init_params(
-    init_function_defs: &Vec<SourceMapped<&Located<StmtKind>>>,
-) -> Result<Vec<Param>, Vec<Error>> {
-    ensure_zero_or_one_function_defs(init_function_defs, CanisterMethodType::Init)?;
+        if post_upgrade_function_defs.len() > 1 {
+            return Err(vec![
+                MultipleSystemMethods::err_from_stmt(
+                    &post_upgrade_function_defs,
+                    CanisterMethodType::PostUpgrade,
+                )
+                .into(),
+                post_upgrade_methods.err().unwrap_or_default(),
+            ]
+            .concat());
+        }
 
-    match init_function_defs.get(0) {
-        Some(init_function_def) => init_function_def.build_params(InternalOrExternal::Internal),
-        None => Ok(vec![]),
+        Ok(match post_upgrade_methods?.pop() {
+            Some(post_upgrade_method) => post_upgrade_method,
+            None => build_post_upgrade_method(None, &self.entry_module_name)?,
+        })
     }
 }
 
-fn build_post_upgrade_info(
+fn build_post_upgrade_methods(
     post_upgrade_function_defs: &Vec<SourceMapped<&Located<StmtKind>>>,
-) -> Result<Vec<Param>, Vec<Error>> {
-    ensure_zero_or_one_function_defs(post_upgrade_function_defs, CanisterMethodType::PostUpgrade)?;
-
-    match post_upgrade_function_defs.get(0) {
-        Some(post_upgrade_function_def) => build_post_upgrade_params(post_upgrade_function_def),
-        None => Ok(vec![]),
-    }
-}
-
-fn build_post_upgrade_params(
-    post_upgrade_function_def: &SourceMapped<&Located<StmtKind>>,
-) -> Result<Vec<Param>, Vec<Error>> {
-    let (post_upgrade_params, _) = (
-        post_upgrade_function_def.build_params(InternalOrExternal::Internal),
-        validate_post_upgrade_return_type(post_upgrade_function_def),
-    )
-        .collect_results()?;
-
-    Ok(post_upgrade_params)
-}
-
-fn select_post_upgrade_method_params(
-    init_params: Vec<Param>,
-    post_upgrade_params: Vec<Param>,
-) -> Vec<Param> {
-    if post_upgrade_params.is_empty() {
-        init_params
-    } else {
-        post_upgrade_params
-    }
-}
-
-fn generate_post_upgrade_method_body(
-    entry_module_name: &String,
-    post_upgrade_function_def: Option<&SourceMapped<&Located<StmtKind>>>,
-) -> Result<TokenStream, Vec<Error>> {
-    rust::generate(post_upgrade_function_def, entry_module_name)
+    entry_module_name: &str,
+) -> Result<Vec<PostUpgradeMethod>, Vec<Error>> {
+    post_upgrade_function_defs
+        .iter()
+        .map(|post_upgrade_function_def| {
+            build_post_upgrade_method(Some(post_upgrade_function_def), entry_module_name)
+        })
+        .collect_results()
 }
 
-fn ensure_zero_or_one_function_defs(
-    function_defs: &Vec<SourceMapped<&Located<StmtKind>>>,
-    canister_method_type: CanisterMethodType,
-) -> Result<(), Vec<Error>> {
-    if function_defs.len() > 1 {
-        Err(MultipleSystemMethods::err_from_stmt(&function_defs, canister_method_type).into())
-    } else {
-        Ok(())
-    }
-}
+fn build_post_upgrade_method(
+    post_upgrade_function_def_option: Option<&SourceMapped<&Located<StmtKind>>>,
+    entry_module_name: &str,
+) -> Result<PostUpgradeMethod, Vec<Error>> {
+    match post_upgrade_function_def_option {
+        Some(post_upgrade_function_def) => {
+            let (params, return_type) = (
+                post_upgrade_function_def.build_params(InternalOrExternal::Internal),
+                post_upgrade_function_def.build_return_type(),
+            )
+                .collect_results()?;
 
-fn validate_post_upgrade_return_type(
-    post_upgrade_function_def: &SourceMapped<&Located<StmtKind>>,
-) -> Result<(), Vec<Error>> {
-    let return_type = post_upgrade_function_def.build_return_type()?;
+            if !canister_method::is_void(return_type) {
+                return Err(ReturnTypeMustBeVoid::err_from_stmt(
+                    post_upgrade_function_def,
+                    CanisterMethodType::PostUpgrade,
+                )
+                .into());
+            }
 
-    if !canister_method::is_void(return_type) {
-        Err(ReturnTypeMustBeVoid::err_from_stmt(
-            post_upgrade_function_def,
-            CanisterMethodType::PostUpgrade,
-        )
-        .into())
-    } else {
-        Ok(())
+            Ok(PostUpgradeMethod {
+                params: params.clone(),
+                body: rust::generate(post_upgrade_function_def_option, entry_module_name)?,
+            })
+        }
+        None => Ok(PostUpgradeMethod {
+            params: vec![],
+            body: rust::generate(post_upgrade_function_def_option, entry_module_name)?,
+        }),
     }
 }
diff --git a/kybra/compiler/kybra_generate/src/canister_method/post_upgrade_method/mod.rs b/kybra/compiler/kybra_generate/src/canister_method/post_upgrade_method/mod.rs
@@ -1,17 +1,2 @@
-use proc_macro2::TokenStream;
-use quote::quote;
-use rustpython_parser::ast::{Located, StmtKind};
-
-use crate::{source_map::SourceMapped, Error};
-
 pub mod build;
 pub mod rust;
-
-pub fn generate_call(
-    function_def_option: &Option<&SourceMapped<&Located<StmtKind>>>,
-) -> Result<TokenStream, Vec<Error>> {
-    match &function_def_option {
-        Some(function_def) => function_def.generate_call_to_py_function(),
-        None => Ok(quote!()),
-    }
-}
diff --git a/kybra/compiler/kybra_generate/src/canister_method/post_upgrade_method/rust.rs b/kybra/compiler/kybra_generate/src/canister_method/post_upgrade_method/rust.rs
@@ -4,12 +4,11 @@ use proc_macro2::TokenStream;
 use quote::quote;
 use rustpython_parser::ast::{Located, StmtKind};
 
-use super::generate_call;
-use crate::{source_map::SourceMapped, Error};
+use crate::{canister_method::init_method::rust::generate_call, source_map::SourceMapped, Error};
 
 pub fn generate(
     post_upgrade_function_def_option: Option<&SourceMapped<&Located<StmtKind>>>,
-    entry_module_name: &String,
+    entry_module_name: &str,
 ) -> Result<TokenStream, Vec<Error>> {
     let call_to_post_upgrade_py_function = generate_call(&post_upgrade_function_def_option)?;