From 3734ca1c381ae8c00116ee3684c087e992bc4914 Mon Sep 17 00:00:00 2001
From: Akihiro Suda <akihiro.suda.cz@hco.ntt.co.jp>
Date: Mon, 30 May 2022 17:14:40 +0900
Subject: [PATCH] llb: add FileOp optimizer

e.g., llb.Scratch().File(llb.Copy(llb.Git(), "/", "/")) --> llb.Git()

Signed-off-by: Akihiro Suda <akihiro.suda.cz@hco.ntt.co.jp>
---
 client/llb/fileop_test.go                     | 53 +++++++++++++++++++
 client/llb/state.go                           | 33 +++++++++++-
 frontend/dockerfile/dockerfile2llb/convert.go |  3 +-
 3 files changed, 87 insertions(+), 2 deletions(-)

diff --git a/client/llb/fileop_test.go b/client/llb/fileop_test.go
index 9fc511d45f28..b08372e8ffbc 100644
--- a/client/llb/fileop_test.go
+++ b/client/llb/fileop_test.go
@@ -641,6 +641,59 @@ func TestFileCreatedTime(t *testing.T) {
 	require.Equal(t, dt3.UnixNano(), copy.Timestamp)
 }
 
+func TestFileOptimize(t *testing.T) {
+	t.Parallel()
+
+	t.Run("Optimizable", func(t *testing.T) {
+		a := Copy(Git("https://github.com/moby/buildkit.git", "v0.4.2"), "/", "/")
+
+		t.Run("WithOptimization", func(t *testing.T) {
+			st, optimized := Scratch().FileOptimize(a, true)
+			require.Equal(t, true, optimized)
+			def, err := st.Marshal(context.TODO())
+			require.NoError(t, err)
+			m, arr := parseDef(t, def.Def)
+			logParsedDef(t, m, arr)
+			require.Equal(t, 2, len(arr))
+			require.Equal(t, "git://github.com/moby/buildkit.git#v0.4.2", m[arr[1].Inputs[0].Digest].Op.(*pb.Op_Source).Source.Identifier)
+		})
+
+		t.Run("WithoutOptimization", func(t *testing.T) {
+			st, optimized := Scratch().FileOptimize(a, false)
+			require.Equal(t, false, optimized)
+			def, err := st.Marshal(context.TODO())
+			require.NoError(t, err)
+			m, arr := parseDef(t, def.Def)
+			logParsedDef(t, m, arr)
+			require.Equal(t, 3, len(arr))
+			require.IsType(t, "git://github.com/moby/buildkit.git#v0.4.2", m[arr[1].Inputs[0].Digest].Op.(*pb.Op_Source).Source.Identifier)
+			require.IsType(t, &pb.FileAction_Copy{}, m[arr[2].Inputs[0].Digest].Op.(*pb.Op_File).File.Actions[0].Action)
+		})
+	})
+
+	t.Run("NotOptimizable", func(t *testing.T) {
+		a := Copy(Git("https://github.com/moby/buildkit.git", "v0.4.2"), "/", "/foo")
+		st, optimized := Scratch().FileOptimize(a, true)
+		require.Equal(t, false, optimized)
+		def, err := st.Marshal(context.TODO())
+		require.NoError(t, err)
+		m, arr := parseDef(t, def.Def)
+		logParsedDef(t, m, arr)
+		require.Equal(t, 3, len(arr))
+		require.Equal(t, "git://github.com/moby/buildkit.git#v0.4.2", m[arr[1].Inputs[0].Digest].Op.(*pb.Op_Source).Source.Identifier)
+		require.IsType(t, &pb.FileAction_Copy{}, m[arr[2].Inputs[0].Digest].Op.(*pb.Op_File).File.Actions[0].Action)
+	})
+}
+
+func logParsedDef(t testing.TB, m map[digest.Digest]pb.Op, arr []pb.Op) {
+	for i, f := range arr {
+		for j, input := range f.Inputs {
+			op := m[input.Digest]
+			t.Logf("Ops[%d].Inputs[%d].Op=%+v", i, j, op.Op)
+		}
+	}
+}
+
 func parseDef(t *testing.T, def [][]byte) (map[digest.Digest]pb.Op, []pb.Op) {
 	m := map[digest.Digest]pb.Op{}
 	arr := make([]pb.Op, 0, len(def))
diff --git a/client/llb/state.go b/client/llb/state.go
index 7ddc7a5f0f32..04604ece45f2 100644
--- a/client/llb/state.go
+++ b/client/llb/state.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net"
+	"path"
 	"strings"
 
 	"github.com/containerd/containerd/platforms"
@@ -274,12 +275,42 @@ func (s State) Run(ro ...RunOption) ExecState {
 }
 
 func (s State) File(a *FileAction, opts ...ConstraintsOpt) State {
+	res, _ := s.FileOptimize(a, false, opts...)
+	return res
+}
+
+// FileOptimize supports optimization.
+// The optimized result may not contain FileOp.
+//
+// e.g., llb.Scratch().File(llb.Copy(llb.Git(), "/", "/")) --> llb.Git()
+func (s State) FileOptimize(a *FileAction, attemptOptimization bool, opts ...ConstraintsOpt) (State, bool) {
 	var c Constraints
 	for _, o := range opts {
 		o.SetConstraintsOption(&c)
 	}
 
-	return s.WithOutput(NewFileOp(s, a, c).Output())
+	if attemptOptimization {
+		if cpA, ok := a.action.(*fileActionCopy); ok &&
+			a.prev == nil &&
+			path.Clean(cpA.src) == "/" && path.Clean(cpA.dest) == "/" &&
+			cpA.fas == nil &&
+			cpA.info.ChownOpt == nil &&
+			cpA.state != nil {
+			if o := cpA.state.Output(); o != nil {
+				switch o.(type) {
+				case *output:
+					switch o.(*output).vertex.(type) {
+					case *SourceOp:
+						newState := s.WithOutput(o)
+						newState.SetMarshalDefaults(opts...)
+						return newState, true
+					}
+				}
+			}
+		}
+	}
+
+	return s.WithOutput(NewFileOp(s, a, c).Output()), false
 }
 
 func (s State) AddEnv(key, value string) State {
diff --git a/frontend/dockerfile/dockerfile2llb/convert.go b/frontend/dockerfile/dockerfile2llb/convert.go
index fcf0d22e0ba6..bad2c30b98f3 100644
--- a/frontend/dockerfile/dockerfile2llb/convert.go
+++ b/frontend/dockerfile/dockerfile2llb/convert.go
@@ -1141,7 +1141,8 @@ func dispatchCopy(d *dispatchState, cfg copyConfig) error {
 
 		d.state = d.state.WithOutput(llb.Merge([]llb.State{d.state, llb.Scratch().File(a, copyOpts...)}, mergeOpts...).Output())
 	} else {
-		d.state = d.state.File(a, fileOpt...)
+		const attemptOptimization = true
+		d.state, _ = d.state.FileOptimize(a, attemptOptimization, fileOpt...)
 	}
 
 	return commitToHistory(&d.image, commitMessage.String(), true, &d.state)