refactor: move export generation to codegen

cmd/gravity/src/codegen/bindings.rs

@@ -1,30 +1,46 @@
+use std::collections::BTreeMap;
+
 use genco::{prelude::*, tokens::Tokens};
-use wit_bindgen_core::wit_parser::{Record, Resolve, Type, TypeDef, TypeDefKind};
+use wit_bindgen_core::wit_parser::{Resolve, SizeAlign, World};
 
 use crate::{
-    codegen::wasm::{Wasm, WasmData},
-    go::*,
-    resolve_type,
+    codegen::{
+        ExportGenerator, FactoryGenerator,
+        exports::ExportConfig,
+        factory::FactoryConfig,
+        imports::{ImportAnalyzer, ImportCodeGenerator},
+        ir::AnalyzedImports,
+        wasm::{Wasm, WasmData},
+    },
+    go::GoIdentifier,
 };
 
 /// The WIT bindings for a world.
-pub struct Bindings {
+pub struct Bindings<'a> {
+    resolve: &'a Resolve,
+    world: &'a World,
     /// The cumulative output tokens containing the Go bindings.
     // TODO(#16): Don't use the internal bindings.out field
     pub out: Tokens<Go>,
 
     /// The identifier of the Go variable containing the WebAssembly bytes.
     raw_wasm_var: GoIdentifier,
+
+    /// The sizes of the architecture.
+    sizes: &'a SizeAlign,
 }
 
-impl Bindings {
+impl<'a> Bindings<'a> {
     /// Creates a new bindings generator for the selected world.
-    pub fn new(world: &str) -> Self {
-        let wasm_var = GoIdentifier::private(format!("wasm-file-{world}"));
+    pub fn new(resolve: &'a Resolve, world: &'a World, sizes: &'a SizeAlign) -> Self {
+        let world_name = &world.name;
+        let wasm_var = GoIdentifier::private(format!("wasm-file-{world_name}"));
         Self {
-            // world,
+            resolve,
+            world,
             out: Tokens::new(),
             raw_wasm_var: wasm_var,
+            sizes,
         }
     }
 
@@ -33,96 +49,53 @@ impl Bindings {
         Wasm::new(&self.raw_wasm_var, wasm).format_into(&mut self.out)
     }
 
-    pub fn define_type(&mut self, typ_def: &TypeDef, resolve: &Resolve) {
-        let TypeDef { name, kind, .. } = typ_def;
-        match kind {
-            TypeDefKind::Record(Record { fields }) => {
-                let name = GoIdentifier::public(name.as_deref().expect("record to have a name"));
-                let fields = fields.iter().map(|field| {
-                    (
-                        GoIdentifier::public(&field.name),
-                        resolve_type(&field.ty, resolve),
-                    )
-                });
-
-                quote_in! { self.out =>
-                    $['\n']
-                    type $name struct {
-                        $(for (name, typ) in fields join ($['\r']) => $name $typ)
-                    }
-                }
-            }
-            TypeDefKind::Resource => todo!("TODO(#5): implement resources"),
-            TypeDefKind::Handle(_) => todo!("TODO(#5): implement resources"),
-            TypeDefKind::Flags(_) => todo!("TODO(#4):generate flags type definition"),
-            TypeDefKind::Tuple(_) => todo!("TODO(#4):generate tuple type definition"),
-            TypeDefKind::Variant(_) => {
-                // TODO(#4): Generate aliases if the variant name doesn't match the struct name
-            }
-            TypeDefKind::Enum(inner) => {
-                let name = name.clone().expect("enum to have a name");
-                let enum_type = &GoIdentifier::private(&name);
-
-                let enum_interface = GoIdentifier::public(&name);
-
-                let enum_function = &GoIdentifier::private(format!("is-{}", &name));
-
-                let variants = inner
-                    .cases
-                    .iter()
-                    .map(|variant| GoIdentifier::public(&variant.name));
+    /// Generate the bindings.
+    ///
+    /// This generates the imports (interfaces, types, functions), the factory and instance
+    /// type, and the exports (functions).
+    pub fn generate(&mut self) {
+        let (imports, chains) = self.generate_imports();
+        self.generate_factory(&imports, chains);
+        self.generate_exports(&imports.instance_name);
+    }
 
-                quote_in! { self.out =>
-                    $['\n']
-                    type $enum_interface interface {
-                        $enum_function()
-                    }
+    /// Generates the imports for the bindings.
+    fn generate_imports(&mut self) -> (AnalyzedImports, BTreeMap<String, Tokens<Go>>) {
+        let analyzer = ImportAnalyzer::new(self.resolve, self.world);
+        let analyzed = analyzer.analyze();
 
-                    type $enum_type int
+        let generator = ImportCodeGenerator::new(self.resolve, &analyzed, self.sizes);
+        let import_chains = generator.import_chains();
+        generator.format_into(&mut self.out);
+        (analyzed, import_chains)
+    }
 
-                    func ($enum_type) $enum_function() {}
+    /// Generates the factory and instantiate functions, including any
+    /// required interfaces.
+    fn generate_factory(
+        &mut self,
+        analyzed_imports: &AnalyzedImports,
+        import_chains: BTreeMap<String, Tokens<Go>>,
+    ) {
+        let config = FactoryConfig {
+            analyzed_imports,
+            import_chains,
+            wasm_var_name: &self.raw_wasm_var,
+        };
+        FactoryGenerator::new(config).format_into(&mut self.out)
+    }
 
-                    const (
-                        $(for name in variants join ($['\r']) => $name $enum_type = iota)
-                    )
-                }
-            }
-            TypeDefKind::Option(_) => todo!("TODO(#4): generate option type definition"),
-            TypeDefKind::Result(_) => todo!("TODO(#4): generate result type definition"),
-            TypeDefKind::List(_) => todo!("TODO(#4): generate list type definition"),
-            TypeDefKind::Future(_) => todo!("TODO(#4): generate future type definition"),
-            TypeDefKind::Stream(_) => todo!("TODO(#4): generate stream type definition"),
-            TypeDefKind::Type(Type::Id(_)) => {
-                // TODO(#4):  Only skip this if we have already generated the type
-            }
-            TypeDefKind::Type(Type::Bool) => todo!("TODO(#4): generate bool type alias"),
-            TypeDefKind::Type(Type::U8) => todo!("TODO(#4): generate u8 type alias"),
-            TypeDefKind::Type(Type::U16) => todo!("TODO(#4): generate u16 type alias"),
-            TypeDefKind::Type(Type::U32) => todo!("TODO(#4): generate u32 type alias"),
-            TypeDefKind::Type(Type::U64) => todo!("TODO(#4): generate u64 type alias"),
-            TypeDefKind::Type(Type::S8) => todo!("TODO(#4): generate s8 type alias"),
-            TypeDefKind::Type(Type::S16) => todo!("TODO(#4): generate s16 type alias"),
-            TypeDefKind::Type(Type::S32) => todo!("TODO(#4): generate s32 type alias"),
-            TypeDefKind::Type(Type::S64) => todo!("TODO(#4): generate s64 type alias"),
-            TypeDefKind::Type(Type::F32) => todo!("TODO(#4): generate f32 type alias"),
-            TypeDefKind::Type(Type::F64) => todo!("TODO(#4): generate f64 type alias"),
-            TypeDefKind::Type(Type::Char) => todo!("TODO(#4): generate char type alias"),
-            TypeDefKind::Type(Type::String) => {
-                let name =
-                    GoIdentifier::public(name.as_deref().expect("string alias to have a name"));
-                // TODO(#4): We might want a Type Definition (newtype) instead of Type Alias here
-                quote_in! { self.out =>
-                    $['\n']
-                    type $name = string
-                }
-            }
-            TypeDefKind::Type(Type::ErrorContext) => {
-                todo!("TODO(#4): generate error context definition")
-            }
-            TypeDefKind::FixedSizeList(_, _) => {
-                todo!("TODO(#4): generate fixed size list definition")
-            }
-            TypeDefKind::Unknown => panic!("cannot generate Unknown type"),
-        }
+    /// Generates all exports for the world.
+    ///
+    /// Note: for now this only generates functions; types and interfaces are
+    /// still TODO
+    fn generate_exports(&mut self, instance: &GoIdentifier) {
+        let config = ExportConfig {
+            instance,
+            world: self.world,
+            resolve: self.resolve,
+            sizes: self.sizes,
+        };
+        ExportGenerator::new(config).format_into(&mut self.out)
     }
 }

cmd/gravity/src/codegen/exports.rs

@@ -0,0 +1,172 @@
+use genco::prelude::*;
+use wit_bindgen_core::wit_parser::{Function, Resolve, SizeAlign, World, WorldItem};
+
+use crate::go::{GoIdentifier, GoResult, GoType, imports::CONTEXT_CONTEXT};
+
+pub struct ExportConfig<'a> {
+    pub instance: &'a GoIdentifier,
+    pub world: &'a World,
+    pub resolve: &'a Resolve,
+    pub sizes: &'a SizeAlign,
+}
+
+pub struct ExportGenerator<'a> {
+    config: ExportConfig<'a>,
+}
+
+impl<'a> ExportGenerator<'a> {
+    pub fn new(config: ExportConfig<'a>) -> Self {
+        Self { config }
+    }
+
+    /// Generate the Go function code for the given function.
+    ///
+    /// The signature is obtained by:
+    /// - getting the function parameters from the `wit_parser::Function`, converting
+    ///   names to to Go identifiers and types to Go types.
+    /// - similar for the result
+    ///
+    /// To implement the body, we:
+    /// - creating a `Func` struct which implements `Bindgen` and passing it to the
+    ///   `wit_bindgen_core::abi::call` function. This will call `Func::emit` lots of
+    ///   times, one for each instruction in the function, and `Func::emit` will generate
+    ///   Go code for each instruction
+    fn generate_function(&self, func: &Function, tokens: &mut Tokens<Go>) {
+        let params = func
+            .params
+            .iter()
+            .map(
+                |(name, wit_type)| match crate::resolve_type(wit_type, self.config.resolve) {
+                    GoType::ValueOrOk(t) => (GoIdentifier::local(name), *t),
+                    t => (GoIdentifier::local(name), t),
+                },
+            )
+            .collect::<Vec<_>>();
+
+        let result = if let Some(wit_type) = &func.result {
+            GoResult::Anon(crate::resolve_type(wit_type, self.config.resolve))
+        } else {
+            GoResult::Empty
+        };
+
+        let mut f = crate::Func::export(result, self.config.sizes);
+        wit_bindgen_core::abi::call(
+            self.config.resolve,
+            wit_bindgen_core::abi::AbiVariant::GuestExport,
+            wit_bindgen_core::abi::LiftLower::LowerArgsLiftResults,
+            func,
+            &mut f,
+            // async is not currently supported
+            false,
+        );
+
+        let arg_assignments = f
+            .args()
+            .iter()
+            .zip(&params)
+            .map(|(arg, (param, _))| (arg, param))
+            .collect::<Vec<_>>();
+        let fn_name = &GoIdentifier::public(&func.name);
+        quote_in! { *tokens =>
+            $['\n']
+            func (i *$(self.config.instance)) $fn_name(
+                $['\r']
+                ctx $CONTEXT_CONTEXT,
+                $(for (name, typ) in &params join ($['\r']) => $name $typ,)
+            ) $(f.result()) {
+                $(for (arg, param) in arg_assignments join ($['\r']) => $arg := $param)
+                $(f.body())
+            }
+        }
+    }
+}
+
+impl FormatInto<Go> for ExportGenerator<'_> {
+    fn format_into(self, tokens: &mut Tokens<Go>) {
+        for item in self.config.world.exports.values() {
+            match item {
+                WorldItem::Function(func) => self.generate_function(func, tokens),
+                WorldItem::Interface { .. } => todo!("generate interface exports"),
+                WorldItem::Type(_) => todo!("generate type exports"),
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use genco::prelude::*;
+    use wit_bindgen_core::wit_parser::{
+        Function, FunctionKind, Resolve, SizeAlign, Type, World, WorldItem, WorldKey,
+    };
+
+    use crate::go::GoIdentifier;
+
+    use super::{ExportConfig, ExportGenerator};
+
+    #[test]
+    fn test_generate_function_simple_u32_param() {
+        let func = Function {
+            name: "add_number".to_string(),
+            kind: FunctionKind::Freestanding,
+            params: vec![("value".to_string(), Type::U32)],
+            result: Some(Type::U32),
+            docs: Default::default(),
+            stability: Default::default(),
+        };
+
+        let world = World {
+            name: "test-world".to_string(),
+            imports: [].into(),
+            exports: [(
+                WorldKey::Name("add-number".to_string()),
+                WorldItem::Function(func.clone()),
+            )]
+            .into(),
+            docs: Default::default(),
+            stability: Default::default(),
+            includes: Default::default(),
+            include_names: Default::default(),
+            package: None,
+        };
+
+        let resolve = Resolve::new();
+        let mut sizes = SizeAlign::default();
+        sizes.fill(&resolve);
+        let instance = GoIdentifier::public("TestInstance");
+
+        let config = ExportConfig {
+            instance: &instance,
+            world: &world,
+            resolve: &resolve,
+            sizes: &sizes,
+        };
+
+        let generator = ExportGenerator::new(config);
+        let mut tokens = Tokens::new();
+
+        // Call the actual generate_function method
+        generator.generate_function(&func, &mut tokens);
+
+        let generated = tokens.to_string().unwrap();
+        println!("Generated: {}", generated);
+
+        // Verify basic function structure
+        assert!(generated.contains("func (i *TestInstance) AddNumber("));
+        assert!(generated.contains("value uint32"));
+        assert!(generated.contains("ctx context.Context"));
+        assert!(generated.contains(") uint32 {"));
+
+        // Verify function body
+        assert!(generated.contains("arg0 := value"));
+        assert!(
+            generated
+                .contains("i.module.ExportedFunction(\"add_number\").Call(ctx, uint64(result0))")
+        );
+        assert!(generated.contains("if err1 != nil {"));
+        assert!(generated.contains("panic(err1)"));
+        assert!(generated.contains("results1 := raw1[0]"));
+        assert!(generated.contains("result2 := api.DecodeU32(results1)"));
+        assert!(generated.contains("return result2"));
+    }
+}