randomness figured out, asynchronously after init/post_upgrade

demergent-labs · Apr 1, 2024 · 679e6f5 · 679e6f5
1 parent 48c133c
commit 679e6f5
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 4 deletions.
diff --git a/examples/randomness/src/main.did b/examples/randomness/src/main.did
@@ -1 +1,4 @@
-service : () -> { random_number : () -> (float64) }
+service : () -> {
+  does_interpreter_exist : () -> (bool) query;
+  random_number : () -> (float64);
+}
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
@@ -6,7 +6,7 @@ use rustpython_parser::ast::{Located, StmtKind};
 use crate::canister_method::post_upgrade_method::generate_call;
 use crate::canister_method::post_upgrade_method::rust::{
     generate_call_to_user_init_or_post_upgrade, generate_code_init, generate_ic_object_init,
-    generate_interpreter_init, generate_save_global_interpreter,
+    generate_interpreter_init, generate_randomness, generate_save_global_interpreter,
 };
 use crate::{source_map::SourceMapped, Error};
 
@@ -29,9 +29,10 @@ pub fn generate(
         &call_to_init_py_function,
         &call_to_post_upgrade_py_function,
     );
+    let randomness = generate_randomness();
 
     Ok(quote! {
-        unsafe { ic_wasi_polyfill::init(&[], &[]); };
+        ic_wasi_polyfill::init(&[], &[]);
 
         #interpreter_init
 
@@ -42,5 +43,7 @@ pub fn generate(
         #save_global_interpreter
 
         #call_to_user_init_or_post_upgrade
+
+        #randomness
     })
 }
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
@@ -27,9 +27,10 @@ pub fn generate(
         &call_to_init_py_function,
         &call_to_post_upgrade_py_function,
     );
+    let randomness = generate_randomness();
 
     Ok(quote! {
-        unsafe { ic_wasi_polyfill::init(&[], &[]); };
+        ic_wasi_polyfill::init(&[], &[]);
 
         #interpreter_init
 
@@ -40,6 +41,8 @@ pub fn generate(
         #save_global_interpreter
 
         #call_to_user_init_or_post_upgrade
+
+        #randomness
     })
 }
 
@@ -152,3 +155,30 @@ pub fn generate_call_to_user_init_or_post_upgrade(
             CANISTER_INITIALIZED_REF_CELL.with(|canister_initialized_ref_cell| canister_initialized_ref_cell.borrow_mut().set(1).unwrap_or_trap());
     }
 }
+
+pub fn generate_randomness() -> TokenStream {
+    quote! {
+        ic_cdk_timers::set_timer(std::time::Duration::from_secs(0), || {
+            ic_cdk::spawn(async move {
+                let result: ic_cdk::api::call::CallResult<(Vec<u8>,)> = ic_cdk::api::management_canister::main::raw_rand().await;
+
+                match result {
+                    Ok(randomness) => {
+                        let interpreter = unsafe { INTERPRETER_OPTION.as_mut() }
+                            .ok_or_else(|| "SystemError: missing python interpreter".to_string()).unwrap();
+                        let scope = unsafe { SCOPE_OPTION.as_mut() }
+                            .ok_or_else(|| "SystemError: missing python scope".to_string()).unwrap();
+
+                        interpreter.enter(|vm| {
+                            let random_module = vm.import("random", None, 0).unwrap();
+                            let seed_fn = random_module.get_attr("seed", vm).unwrap();
+
+                            seed_fn.call((vm.ctx.new_bytes(randomness.0),), vm).unwrap();
+                        });
+                    },
+                    Err(err) => panic!(err)
+                };
+            });
+        });
+    }
+}