From 1601dd2b005b59d1d0503afcf77e84385c7ff01f Mon Sep 17 00:00:00 2001
From: Norman Meier <norman@samourai.coop>
Date: Sun, 30 Jun 2024 01:11:41 +0200
Subject: [PATCH] feat(transpile): allow extra root dirs

Signed-off-by: Norman Meier <norman@samourai.coop>
---
 gnovm/cmd/gno/transpile.go              | 10 +++++-
 gnovm/pkg/transpiler/transpiler.go      | 46 +++++++++++++++++++++++--
 gnovm/pkg/transpiler/transpiler_test.go |  2 +-
 3 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/gnovm/cmd/gno/transpile.go b/gnovm/cmd/gno/transpile.go
index 2e12ee6f4b3..156d6823636 100644
--- a/gnovm/cmd/gno/transpile.go
+++ b/gnovm/cmd/gno/transpile.go
@@ -25,6 +25,7 @@ import (
 type transpileCfg struct {
 	verbose     bool
 	rootDir     string
+	extraDirs   string
 	skipImports bool
 	gobuild     bool
 	goBinary    string
@@ -94,6 +95,13 @@ func (c *transpileCfg) RegisterFlags(fs *flag.FlagSet) {
 		"clone location of github.com/gnolang/gno (gno tries to guess it)",
 	)
 
+	fs.StringVar(
+		&c.extraDirs,
+		"extra-dirs",
+		"",
+		"extra directories to look for packages in",
+	)
+
 	fs.BoolVar(
 		&c.skipImports,
 		"skip-imports",
@@ -251,7 +259,7 @@ func transpileFile(srcPath string, opts *transpileOptions) error {
 	targetFilename, tags := transpiler.TranspiledFilenameAndTags(srcPath)
 
 	// preprocess.
-	transpileRes, err := transpiler.Transpile(string(source), tags, srcPath)
+	transpileRes, err := transpiler.Transpile(string(source), tags, srcPath, strings.Split(opts.cfg.extraDirs, ","))
 	if err != nil {
 		return fmt.Errorf("transpile: %w", err)
 	}
diff --git a/gnovm/pkg/transpiler/transpiler.go b/gnovm/pkg/transpiler/transpiler.go
index bd4bb1b1bc9..841c831cfcc 100644
--- a/gnovm/pkg/transpiler/transpiler.go
+++ b/gnovm/pkg/transpiler/transpiler.go
@@ -10,9 +10,11 @@ import (
 	"go/parser"
 	goscanner "go/scanner"
 	"go/token"
+	"io/fs"
 	"os"
 	"path"
 	"path/filepath"
+	"regexp"
 	"strconv"
 	"strings"
 
@@ -75,7 +77,7 @@ func TranspiledFilenameAndTags(gnoFilePath string) (targetFilename, tags string)
 // Transpile performs transpilation on the given source code. tags can be used
 // to specify build tags; and filename helps generate useful error messages and
 // discriminate between test and normal source files.
-func Transpile(source, tags, filename string) (*Result, error) {
+func Transpile(source, tags, filename string, extraDirs []string) (*Result, error) {
 	fset := token.NewFileSet()
 	f, err := parser.ParseFile(fset, filename, source,
 		// SkipObjectResolution -- unused here.
@@ -87,7 +89,8 @@ func Transpile(source, tags, filename string) (*Result, error) {
 
 	isTestFile := strings.HasSuffix(filename, "_test.gno") || strings.HasSuffix(filename, "_filetest.gno")
 	ctx := &transpileCtx{
-		rootDir: gnoenv.RootDir(),
+		rootDir:   gnoenv.RootDir(),
+		extraDirs: extraDirs,
 	}
 	stdlibPrefix := filepath.Join(ctx.rootDir, "gnovm", "stdlibs")
 	if isTestFile {
@@ -142,6 +145,9 @@ type transpileCtx struct {
 	// If rootDir is given, we will check that the directory of the import path
 	// exists (using rootDir/packageDirLocation()).
 	rootDir string
+	// If extraDirs is given, we will check that the directory of the import path
+	// exists (using rootDir/packageDirLocation()).
+	extraDirs []string
 	// This should be set if we're working with a file from a standard library.
 	// This allows us to easily check if a function has a native binding, and as
 	// such modify its call expressions appropriately.
@@ -165,7 +171,41 @@ func (ctx *transpileCtx) transformFile(fset *token.FileSet, f *ast.File) (*ast.F
 				continue
 			}
 
-			if ctx.rootDir != "" {
+			foundInExtra := false
+			for _, extraDir := range ctx.extraDirs {
+				// TODO: find packages once
+				modFiles := []string{}
+				if err := filepath.Walk(extraDir, func(path string, info fs.FileInfo, err error) error {
+					if !strings.HasSuffix(path, "/gno.mod") {
+						return nil
+					}
+					modFiles = append(modFiles, path)
+					return nil
+				}); err != nil {
+					return nil, err
+				}
+				packageNames := []string{}
+				for _, modFile := range modFiles {
+					modData, err := os.ReadFile(modFile)
+					if err != nil {
+						continue
+					}
+					re := regexp.MustCompile(`^module (.+)`)
+					res := re.FindSubmatch(modData)
+					if len(res) != 2 {
+						continue
+					}
+					packageNames = append(packageNames, string(res[1]))
+				}
+
+				for _, pn := range packageNames {
+					if pn == importPath {
+						foundInExtra = true
+					}
+				}
+			}
+
+			if !foundInExtra && ctx.rootDir != "" {
 				dirPath := filepath.Join(ctx.rootDir, PackageDirLocation(importPath))
 				if _, err := os.Stat(dirPath); err != nil {
 					if !os.IsNotExist(err) {
diff --git a/gnovm/pkg/transpiler/transpiler_test.go b/gnovm/pkg/transpiler/transpiler_test.go
index 2a0707f7f79..bd341935cc1 100644
--- a/gnovm/pkg/transpiler/transpiler_test.go
+++ b/gnovm/pkg/transpiler/transpiler_test.go
@@ -420,7 +420,7 @@ func testfunc() {
 			if filename == "" {
 				filename = "foo.gno"
 			}
-			res, err := Transpile(source, c.tags, filename)
+			res, err := Transpile(source, c.tags, filename, nil)
 
 			if c.expectedError != "" {
 				require.EqualError(t, err, c.expectedError)