Repository: openllb/hlb
Pull request #197 (2020-10-23 16:49:28)

Add escape fragments to heredocs
Allow for escaped fragments within heredocs to keep consistency with string fragments.

Example:
```hlb
fs default() {
	mkfile "foo" 0o644 <<~EOM
		Escape this \${PATH}
        EOM
}
```


Patch:
diff --git a/codegen/codegen.go b/codegen/codegen.go
index 4688a1c..4fac67a 100644
--- a/codegen/codegen.go
+++ b/codegen/codegen.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"fmt"
 	"reflect"
+	"strconv"
 	"strings"
 
 	"github.com/lithammer/dedent"
@@ -131,10 +132,15 @@ func (cg *CodeGen) EmitStringLit(ctx context.Context, scope *parser.Scope, str *
 		switch {
 		case f.Escaped != nil:
 			escaped := *f.Escaped
-			if len(escaped) == 2 && escaped[1] == '"' {
-				escaped = strings.TrimPrefix(escaped, `\`)
+			if escaped[1] == '$' {
+				pieces = append(pieces, "$")
+			} else {
+				value, _, _, err := strconv.UnquoteChar(escaped, '"')
+				if err != nil {
+					return err
+				}
+				pieces = append(pieces, string(value))
 			}
-			pieces = append(pieces, escaped)
 		case f.Interpolated != nil:
 			exprRet := NewRegister()
 			err := cg.EmitExpr(ctx, scope, f.Interpolated.Expr, nil, nil, nil, exprRet)
@@ -161,6 +167,13 @@ func (cg *CodeGen) EmitHeredoc(ctx context.Context, scope *parser.Scope, heredoc
 		switch {
 		case f.Whitespace != nil:
 			pieces = append(pieces, *f.Whitespace)
+		case f.Escaped != nil:
+			escaped := *f.Escaped
+			if escaped[1] == '$' {
+				pieces = append(pieces, "$")
+			} else {
+				pieces = append(pieces, escaped)
+			}
 		case f.Interpolated != nil:
 			exprRet := NewRegister()
 			err := cg.EmitExpr(ctx, scope, f.Interpolated.Expr, nil, nil, nil, exprRet)
diff --git a/langserver/server.go b/langserver/server.go
index 13c7585..bdd6756 100644
--- a/langserver/server.go
+++ b/langserver/server.go
@@ -364,7 +364,10 @@ func highlightStringFragment(lines map[int]lsp.SemanticHighlightingTokens, f *pa
 }
 
 func highlightHeredocFragment(lines map[int]lsp.SemanticHighlightingTokens, f *parser.HeredocFragment) {
-	if f.Interpolated != nil {
+	switch {
+	case f.Escaped != nil:
+		highlightNode(lines, f, Comment)
+	case f.Interpolated != nil:
 		highlightExpr(lines, f.Interpolated.Expr)
 	}
 }
diff --git a/parser/cst.go b/parser/cst.go
index 9e36555..1f90c00 100644
--- a/parser/cst.go
+++ b/parser/cst.go
@@ -51,6 +51,7 @@ var (
 		"Heredoc": {
 			{"HeredocEnd", `\b\1\b`, stateful.Pop()},
 			{"Whitespace", `\s+`, nil},
+			{"Escaped", `\\.`, nil},
 			{"Interpolated", `\${`, stateful.Push("Interpolated")},
 			{"Text", `\$|[^\s$]+`, nil},
 		},
@@ -697,6 +698,7 @@ type Heredoc struct {
 type HeredocFragment struct {
 	Mixin
 	Whitespace   *string       `parser:"( @Whitespace"`
+	Escaped      *string       `parser:"| @Escaped"`
 	Interpolated *Interpolated `parser:"| @@"`
 	Text         *string       `parser:"| @(Text | RawText) )"`
 }
diff --git a/parser/unparse.go b/parser/unparse.go
index dc18269..db34dec 100644
--- a/parser/unparse.go
+++ b/parser/unparse.go
@@ -529,6 +529,8 @@ func (hf *HeredocFragment) Unparse(opts ...UnparseOption) string {
 	switch {
 	case hf.Whitespace != nil:
 		return *hf.Whitespace
+	case hf.Escaped != nil:
+		return *hf.Escaped
 	case hf.Interpolated != nil:
 		opts = append(opts, WithNoNewline())
 		return hf.Interpolated.Unparse(opts...)


Test patch:
diff --git a/codegen/codegen_test.go b/codegen/codegen_test.go
index d647ff0..af6747d 100644
--- a/codegen/codegen_test.go
+++ b/codegen/codegen_test.go
@@ -626,6 +626,37 @@ func TestCodeGen(t *testing.T) {
 				llb.Mkfile("foo", 0644, []byte("Hello world")),
 			))
 		},
+	}, {
+		"string escape",
+		[]string{"default"},
+		`
+		fs default() {
+			mkfile "foo" 0o644 "Escape \${PATH} Escape \" Escape \n Escape \\"
+		}
+		`, "",
+		func(ctx context.Context, t *testing.T) solver.Request {
+			return Expect(t, llb.Scratch().File(
+				llb.Mkfile("foo", 0644, []byte("Escape ${PATH} Escape \" Escape \n Escape \\")),
+			))
+		},
+	}, {
+		"heredoc escape",
+		[]string{"default"},
+		`
+		fs default() {
+			mkfile "foo" 0o644 <<~EOM
+				Escape \${PATH}
+				Don't escape \"
+				Don't escape \n
+				Don't escape \\
+		        EOM
+		}
+		`, "",
+		func(ctx context.Context, t *testing.T) solver.Request {
+			return Expect(t, llb.Scratch().File(
+				llb.Mkfile("foo", 0644, []byte(`Escape ${PATH} Don't escape \" Don't escape \n Don't escape \\`)),
+			))
+		},
 	}, {
 		"entitlements",
 		[]string{"default"},