From dedaa190515550d4c02887ae0dd7b2eda031a20f Mon Sep 17 00:00:00 2001
From: Jairus <me@jairus.dev>
Date: Fri, 1 Nov 2024 16:48:11 -0400
Subject: [PATCH] setup go transform

---
 .../modus-go-build/extractor/extractor.go     |  6 ++---
 .../modus-go-build/extractor/functions.go     | 26 +++++++++++++++++++
 .../modus-go-build/extractor/transform.go     |  4 ++-
 .../tools/modus-go-build/metadata/metadata.go |  6 +++++
 4 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/sdk/go/tools/modus-go-build/extractor/extractor.go b/sdk/go/tools/modus-go-build/extractor/extractor.go
index 6d5c44ad4..f8cc781cf 100644
--- a/sdk/go/tools/modus-go-build/extractor/extractor.go
+++ b/sdk/go/tools/modus-go-build/extractor/extractor.go
@@ -29,14 +29,14 @@ func CollectProgramInfo(config *config.Config, meta *metadata.Metadata, wasmFunc
 
 	for name, f := range getExportedFunctions(pkgs) {
 		if _, ok := wasmFunctions.Exports[name]; ok {
-			meta.FnExports[name] = transformFunc(name, f)
+			meta.FnExports[name] = transformFunc(name, f, pkgs)
 			findRequiredTypes(f, requiredTypes)
 		}
 	}
 
 	for name, f := range getImportedFunctions(pkgs) {
 		if _, ok := wasmFunctions.Imports[name]; ok {
-			meta.FnImports[name] = transformFunc(name, f)
+			meta.FnImports[name] = transformFunc(name, f, pkgs)
 			findRequiredTypes(f, requiredTypes)
 		}
 	}
@@ -44,7 +44,7 @@ func CollectProgramInfo(config *config.Config, meta *metadata.Metadata, wasmFunc
 	// proxy imports overwrite regular imports
 	for name, f := range getProxyImportFunctions(pkgs) {
 		if _, ok := meta.FnImports[name]; ok {
-			meta.FnImports[name] = transformFunc(name, f)
+			meta.FnImports[name] = transformFunc(name, f, pkgs)
 			findRequiredTypes(f, requiredTypes)
 		}
 	}
diff --git a/sdk/go/tools/modus-go-build/extractor/functions.go b/sdk/go/tools/modus-go-build/extractor/functions.go
index 182af31a6..0a54d4714 100644
--- a/sdk/go/tools/modus-go-build/extractor/functions.go
+++ b/sdk/go/tools/modus-go-build/extractor/functions.go
@@ -16,6 +16,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/hypermodeinc/modus/sdk/go/tools/modus-go-build/metadata"
 	"golang.org/x/tools/go/packages"
 )
 
@@ -26,6 +27,31 @@ var wellKnownTypes = map[string]bool{
 	"time.Duration": true,
 }
 
+func getFuncDocumentation(pkgs map[string]*packages.Package, fn *types.Func) *metadata.Docs {
+	var comments []string
+
+	for _, pkg := range pkgs {
+		for _, file := range pkg.Syntax {
+			for _, decl := range file.Decls {
+				if fd, ok := decl.(*ast.FuncDecl); ok {
+					if fd.Name.Name == fn.Name() {
+						if fd.Doc != nil {
+							for _, comment := range fd.Doc.List {
+								comments = append(comments, comment.Text)
+							}
+						}
+						return &metadata.Docs{
+							Description: strings.Join(comments, "\n"),
+						}
+					}
+				}
+			}
+		}
+	}
+	return nil
+}
+
+
 func getExportedFunctions(pkgs map[string]*packages.Package) map[string]*types.Func {
 	results := make(map[string]*types.Func)
 	for _, pkg := range pkgs {
diff --git a/sdk/go/tools/modus-go-build/extractor/transform.go b/sdk/go/tools/modus-go-build/extractor/transform.go
index d79d9420f..f871c819f 100644
--- a/sdk/go/tools/modus-go-build/extractor/transform.go
+++ b/sdk/go/tools/modus-go-build/extractor/transform.go
@@ -14,6 +14,7 @@ import (
 
 	"github.com/hypermodeinc/modus/sdk/go/tools/modus-go-build/metadata"
 	"github.com/hypermodeinc/modus/sdk/go/tools/modus-go-build/utils"
+	"golang.org/x/tools/go/packages"
 )
 
 func transformStruct(name string, s *types.Struct) *metadata.TypeDefinition {
@@ -37,7 +38,7 @@ func transformStruct(name string, s *types.Struct) *metadata.TypeDefinition {
 	}
 }
 
-func transformFunc(name string, f *types.Func) *metadata.Function {
+func transformFunc(name string, f *types.Func, pkgs map[string]*packages.Package) *metadata.Function {
 	if f == nil {
 		return nil
 	}
@@ -48,6 +49,7 @@ func transformFunc(name string, f *types.Func) *metadata.Function {
 
 	ret := metadata.Function{
 		Name: name,
+		Docs: getFuncDocumentation(pkgs, f),
 	}
 
 	if params != nil {
diff --git a/sdk/go/tools/modus-go-build/metadata/metadata.go b/sdk/go/tools/modus-go-build/metadata/metadata.go
index 8f11af3da..12aa1125b 100644
--- a/sdk/go/tools/modus-go-build/metadata/metadata.go
+++ b/sdk/go/tools/modus-go-build/metadata/metadata.go
@@ -37,16 +37,22 @@ type Metadata struct {
 	Types     TypeMap     `json:"types,omitempty"`
 }
 
+type Docs struct {
+	Description string `json:"description"`
+}
+
 type Function struct {
 	Name       string       `json:"-"`
 	Parameters []*Parameter `json:"parameters,omitempty"`
 	Results    []*Result    `json:"results,omitempty"`
+	Docs       *Docs        `json:"docs,omitempty"`
 }
 
 type TypeDefinition struct {
 	Id     uint32   `json:"id"`
 	Name   string   `json:"-"`
 	Fields []*Field `json:"fields,omitempty"`
+	Docs       *Docs        `json:"docs,omitempty"`
 }
 
 type Parameter struct {