1 files changed, 10 insertions, 6 deletions
diff --git a/mingling/src/program/setup.rs b/mingling/src/program/setup.rs
index fdf7b04..e81247e 100644
--- a/mingling/src/program/setup.rs
+++ b/mingling/src/program/setup.rs
@@ -1,15 +1,19 @@
-use crate::program::Program;
+use crate::{ProgramCollect, program::Program};
 
 mod basic;
 pub use basic::*;
 
-pub trait ProgramSetup {
-    fn setup(program: &mut Program);
+pub trait ProgramSetup<C: ProgramCollect> {
+    fn setup(&mut self, program: &mut Program<C>);
 }
 
-impl Program {
+impl<C> Program<C>
+where
+    C: ProgramCollect,
+{
     /// Load and execute init logic
-    pub fn with_setup<S: ProgramSetup + 'static>(&mut self, _setup: S) {
-        S::setup(self);
+    pub fn with_setup<S: ProgramSetup<C> + 'static>(&mut self, mut setup: S) -> S {
+        S::setup(&mut setup, self);
+        setup
     }
 }