From eec817d3583d2a982106370069e0034f7b642d5a Mon Sep 17 00:00:00 2001
From: Soumil Paranjpay <soumil.paranjpay@gmail.com>
Date: Tue, 12 Dec 2023 19:20:47 +0530
Subject: [PATCH] fix: refactor vm initialization code into method

Signed-off-by: Soumil Paranjpay <soumil.paranjpay@gmail.com>
---
 transformer/external/wasmtransformer.go | 83 +++++++++++--------------
 1 file changed, 35 insertions(+), 48 deletions(-)

diff --git a/transformer/external/wasmtransformer.go b/transformer/external/wasmtransformer.go
index 83596b50d..225c42501 100644
--- a/transformer/external/wasmtransformer.go
+++ b/transformer/external/wasmtransformer.go
@@ -79,30 +79,7 @@ func (t *WASM) GetConfig() (transformertypes.Transformer, *environment.Environme
 
 // DirectoryDetect runs detect in each sub directory
 func (t *WASM) DirectoryDetect(dir string) (map[string][]transformertypes.Artifact, error) {
-	wasmedge.SetLogErrorLevel()
-	conf := wasmedge.NewConfigure(wasmedge.WASI)
-	vm := wasmedge.NewVMWithConfig(conf)
-
-	wasi := vm.GetImportModule(wasmedge.WASI)
-	wasi.InitWasi(
-		[]string{},
-		t.prepareEnv(),
-		[]string{dir + ":" + dir},
-	)
-
-	err := vm.LoadWasmFile(t.WASMConfig.WASMModule)
-	if err != nil {
-		return nil, err
-	}
-	err = vm.Validate()
-	if err != nil {
-		return nil, err
-	}
-	err = vm.Instantiate()
-	if err != nil {
-		return nil, err
-	}
-	_, err = vm.Execute("_start")
+	vm, err := t.initVm([]string{dir + ":" + dir})
 	if err != nil {
 		return nil, err
 	}
@@ -157,11 +134,6 @@ func (t *WASM) Transform(newArtifacts []transformertypes.Artifact, alreadySeenAr
 	}
 	dataStr := string(dataByt)
 
-	wasmedge.SetLogErrorLevel()
-	conf := wasmedge.NewConfigure(wasmedge.WASI)
-	vm := wasmedge.NewVMWithConfig(conf)
-
-	wasi := vm.GetImportModule(wasmedge.WASI)
 	preopens := []string{}
 	for _, artifact := range newArtifacts {
 		for _, paths := range artifact.Paths {
@@ -199,28 +171,11 @@ func (t *WASM) Transform(newArtifacts []transformertypes.Artifact, alreadySeenAr
 		finalPreopens = append(finalPreopens, path+":"+path)
 	}
 
-	wasi.InitWasi(
-		[]string{},
-		t.prepareEnv(),
-		finalPreopens,
-	)
-
-	err = vm.LoadWasmFile(t.WASMConfig.WASMModule)
-	if err != nil {
-		return nil, nil, err
-	}
-	err = vm.Validate()
-	if err != nil {
-		return nil, nil, err
-	}
-	err = vm.Instantiate()
-	if err != nil {
-		return nil, nil, err
-	}
-	_, err = vm.Execute("_start")
+	vm, err := t.initVm(finalPreopens)
 	if err != nil {
 		return nil, nil, err
 	}
+
 	allocateResult, err := vm.Execute("malloc", int32(len(dataStr)+1))
 	if err != nil {
 		return nil, nil, err
@@ -253,3 +208,35 @@ func (t *WASM) Transform(newArtifacts []transformertypes.Artifact, alreadySeenAr
 	createdArtifacts = append(createdArtifacts, output.CreatedArtifacts...)
 	return pathMappings, createdArtifacts, err
 }
+
+func (t *WASM) initVm(preopens []string) (*wasmedge.VM, error) {
+	wasmedge.SetLogErrorLevel()
+	conf := wasmedge.NewConfigure(wasmedge.WASI)
+	vm := wasmedge.NewVMWithConfig(conf)
+
+	wasi := vm.GetImportModule(wasmedge.WASI)
+	wasi.InitWasi(
+		[]string{},
+		t.prepareEnv(),
+		preopens,
+	)
+
+	err := vm.LoadWasmFile(t.WASMConfig.WASMModule)
+	if err != nil {
+		return nil, err
+	}
+	err = vm.Validate()
+	if err != nil {
+		return nil, err
+	}
+	err = vm.Instantiate()
+	if err != nil {
+		return nil, err
+	}
+	_, err = vm.Execute("_start")
+	if err != nil {
+		return nil, err
+	}
+
+	return vm, nil
+}