From 6d505a2cb9c2fb5f3fb20b41f8638c5853abff43 Mon Sep 17 00:00:00 2001
From: Ce Gao <ce.gao@outlook.com>
Date: Mon, 9 May 2022 12:52:22 +0800
Subject: [PATCH] feat(lang): Support command exec (#119)

Signed-off-by: Ce Gao <cegao@tensorchord.ai>
---
 pkg/lang/frontend/starlark/const.go |  1 +
 pkg/lang/frontend/starlark/rules.go | 24 ++++++++++++++++++++-
 pkg/lang/ir/graph.go                | 33 +++++++++++++++++++++++------
 pkg/lang/ir/interface.go            |  6 ++++++
 pkg/lang/ir/types.go                |  4 +---
 5 files changed, 57 insertions(+), 11 deletions(-)

diff --git a/pkg/lang/frontend/starlark/const.go b/pkg/lang/frontend/starlark/const.go
index b6b9925c1..5c0b67b10 100644
--- a/pkg/lang/frontend/starlark/const.go
+++ b/pkg/lang/frontend/starlark/const.go
@@ -24,4 +24,5 @@ const (
 	rulePyPIMirror    = "pip_mirror"
 	ruleShell         = "shell"
 	ruleJupyter       = "jupyter"
+	ruleRun           = "run"
 )
diff --git a/pkg/lang/frontend/starlark/rules.go b/pkg/lang/frontend/starlark/rules.go
index 2224215b5..073cf1925 100644
--- a/pkg/lang/frontend/starlark/rules.go
+++ b/pkg/lang/frontend/starlark/rules.go
@@ -39,7 +39,7 @@ func registerMIDIRules() {
 	starlark.Universe[rulePyPIMirror] = starlark.NewBuiltin(rulePyPIMirror, ruleFuncPyPIMirror)
 	starlark.Universe[ruleShell] = starlark.NewBuiltin(ruleShell, ruleFuncShell)
 	starlark.Universe[ruleJupyter] = starlark.NewBuiltin(ruleJupyter, ruleFuncJupyter)
-
+	starlark.Universe[ruleRun] = starlark.NewBuiltin(ruleRun, ruleFuncRun)
 }
 
 func ruleFuncBase(thread *starlark.Thread, _ *starlark.Builtin,
@@ -262,3 +262,25 @@ func ruleFuncJupyter(thread *starlark.Thread, _ *starlark.Builtin,
 
 	return starlark.None, nil
 }
+
+func ruleFuncRun(thread *starlark.Thread, _ *starlark.Builtin,
+	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var commands *starlark.List
+
+	if err := starlark.UnpackArgs(rulePyPIPackage,
+		args, kwargs, "commands?", &commands); err != nil {
+		return nil, err
+	}
+
+	goCommands := []string{}
+	if commands != nil {
+		for i := 0; i < commands.Len(); i++ {
+			goCommands = append(goCommands, commands.Index(i).(starlark.String).GoString())
+		}
+	}
+
+	logger.Debugf("rule `%s` is invoked, commands=%v", ruleRun, goCommands)
+	ir.Run(goCommands)
+
+	return starlark.None, nil
+}
diff --git a/pkg/lang/ir/graph.go b/pkg/lang/ir/graph.go
index 154124151..0becc2dc7 100644
--- a/pkg/lang/ir/graph.go
+++ b/pkg/lang/ir/graph.go
@@ -43,7 +43,7 @@ func NewGraph() *Graph {
 
 		PyPIPackages:   []string{},
 		SystemPackages: []string{},
-		Exec:           []llb.State{},
+		Exec:           []string{},
 		Shell:          shellBASH,
 	}
 }
@@ -90,16 +90,21 @@ func (g Graph) Compile() (llb.State, error) {
 	if err != nil {
 		return llb.State{}, errors.Wrap(err, "failed to get vscode plugins")
 	}
+
+	var merged llb.State
 	if vscodeStage != nil {
-		merged := llb.Merge([]llb.State{
+		merged = llb.Merge([]llb.State{
 			builtinSystemStage, systemStage, pypiStage, sshStage, *vscodeStage, diffShellStage,
 		})
-		return merged, nil
+	} else {
+		merged = llb.Merge([]llb.State{
+			builtinSystemStage, systemStage, pypiStage, sshStage, diffShellStage,
+		})
 	}
-	merged := llb.Merge([]llb.State{
-		builtinSystemStage, systemStage, pypiStage, sshStage, diffShellStage,
-	})
-	return merged, nil
+
+	// TODO(gaocegege): Support order-based exec.
+	run := g.compileRun(merged)
+	return run, nil
 }
 
 func (g *Graph) compileBase() llb.State {
@@ -274,3 +279,17 @@ func (g Graph) compileZSH(root llb.State) (llb.State, error) {
 	run := zshStage.Run(llb.Shlex(fmt.Sprintf("bash %s", installPath)))
 	return run.Root(), nil
 }
+
+func (g Graph) compileRun(root llb.State) llb.State {
+	if len(g.Exec) == 0 {
+		return root
+	} else if len(g.Exec) == 1 {
+		return root.Run(llb.Shlex(g.Exec[0])).Root()
+	}
+
+	run := root.Run(llb.Shlex(g.Exec[0]))
+	for _, c := range g.Exec[1:] {
+		run = run.Run(llb.Shlex(c))
+	}
+	return run.Root()
+}
diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go
index 3f1dbd805..bab30910c 100644
--- a/pkg/lang/ir/interface.go
+++ b/pkg/lang/ir/interface.go
@@ -76,3 +76,9 @@ func Jupyter(pwd string, port int64) error {
 	}
 	return nil
 }
+
+func Run(commands []string) error {
+	// TODO(gaocegege): Support order-based exec.
+	DefaultGraph.Exec = commands
+	return nil
+}
diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go
index 5467f1927..2d92ceee8 100644
--- a/pkg/lang/ir/types.go
+++ b/pkg/lang/ir/types.go
@@ -15,8 +15,6 @@
 package ir
 
 import (
-	"github.com/moby/buildkit/client/llb"
-
 	"github.com/tensorchord/MIDI/pkg/vscode"
 )
 
@@ -37,7 +35,7 @@ type Graph struct {
 	SystemPackages        []string
 	VSCodePlugins         []vscode.Plugin
 
-	Exec []llb.State
+	Exec []string
 	*JupyterConfig
 }