From 76a0da3937c59b54401a2cea562d66be2767261a Mon Sep 17 00:00:00 2001
From: Mike Farah <mikefarah@gmail.com>
Date: Tue, 12 Mar 2024 15:57:26 +1100
Subject: [PATCH] Refactoring load op

---
 pkg/yqlib/lexer_participle.go | 16 +++++-------
 pkg/yqlib/operation.go        |  3 ++-
 pkg/yqlib/operator_load.go    | 48 ++++++++++++++++++++++++++---------
 3 files changed, 45 insertions(+), 22 deletions(-)

diff --git a/pkg/yqlib/lexer_participle.go b/pkg/yqlib/lexer_participle.go
index 54cf723c..b7ccc5d2 100644
--- a/pkg/yqlib/lexer_participle.go
+++ b/pkg/yqlib/lexer_participle.go
@@ -87,15 +87,13 @@ var participleYqRules = []*participleYqRule{
 	{"Uri", `@uri`, encodeWithIndent(UriFormat, 0), 0},
 	{"SH", `@sh`, encodeWithIndent(ShFormat, 0), 0},
 
-	{"LoadXML", `load_?xml|xml_?load`, loadOp(NewXMLDecoder(ConfiguredXMLPreferences), false), 0},
+	{"LoadXML", `load_?xml|xml_?load`, loadOp(NewXMLDecoder(ConfiguredXMLPreferences)), 0},
 
-	{"LoadBase64", `load_?base64`, loadOp(NewBase64Decoder(), false), 0},
+	{"LoadBase64", `load_?base64`, loadOp(NewBase64Decoder()), 0},
 
-	{"LoadProperties", `load_?props`, loadOp(NewPropertiesDecoder(), false), 0},
-
-	{"LoadString", `load_?str|str_?load`, loadOp(nil, true), 0},
-
-	{"LoadYaml", `load`, loadOp(NewYamlDecoder(LoadYamlPreferences), false), 0},
+	{"LoadProperties", `load_?props`, loadOp(NewPropertiesDecoder()), 0},
+	simpleOp("load_?str|str_?load", loadStringOpType),
+	{"LoadYaml", `load`, loadOp(NewYamlDecoder(LoadYamlPreferences)), 0},
 
 	{"SplitDocument", `splitDoc|split_?doc`, opToken(splitDocumentOpType), 0},
 
@@ -549,8 +547,8 @@ func decodeOp(format *Format) yqAction {
 	return opTokenWithPrefs(decodeOpType, nil, prefs)
 }
 
-func loadOp(decoder Decoder, loadAsString bool) yqAction {
-	prefs := loadPrefs{decoder: decoder, loadAsString: loadAsString}
+func loadOp(decoder Decoder) yqAction {
+	prefs := loadPrefs{decoder}
 	return opTokenWithPrefs(loadOpType, nil, prefs)
 }
 
diff --git a/pkg/yqlib/operation.go b/pkg/yqlib/operation.go
index 2350a4c1..ea29e10d 100644
--- a/pkg/yqlib/operation.go
+++ b/pkg/yqlib/operation.go
@@ -155,7 +155,8 @@ var trimOpType = &operationType{Type: "TRIM", NumArgs: 0, Precedence: 50, Handle
 var toStringOpType = &operationType{Type: "TO_STRING", NumArgs: 0, Precedence: 50, Handler: toStringOperator}
 var stringInterpolationOpType = &operationType{Type: "STRING_INT", NumArgs: 0, Precedence: 50, Handler: stringInterpolationOperator, ToString: valueToStringFunc}
 
-var loadOpType = &operationType{Type: "LOAD", NumArgs: 1, Precedence: 52, Handler: loadYamlOperator}
+var loadOpType = &operationType{Type: "LOAD", NumArgs: 1, Precedence: 52, Handler: loadOperator}
+var loadStringOpType = &operationType{Type: "LOAD_STRING", NumArgs: 1, Precedence: 52, Handler: loadStringOperator}
 
 var keysOpType = &operationType{Type: "KEYS", NumArgs: 0, Precedence: 50, Handler: keysOperator}
 
diff --git a/pkg/yqlib/operator_load.go b/pkg/yqlib/operator_load.go
index 8f0a5003..f30ee352 100644
--- a/pkg/yqlib/operator_load.go
+++ b/pkg/yqlib/operator_load.go
@@ -15,8 +15,7 @@ var LoadYamlPreferences = YamlPreferences{
 }
 
 type loadPrefs struct {
-	loadAsString bool
-	decoder      Decoder
+	decoder Decoder
 }
 
 func loadString(filename string) (*CandidateNode, error) {
@@ -31,7 +30,7 @@ func loadString(filename string) (*CandidateNode, error) {
 	return &CandidateNode{Kind: ScalarNode, Tag: "!!str", Value: string(filebytes)}, nil
 }
 
-func loadYaml(filename string, decoder Decoder) (*CandidateNode, error) {
+func loadWithDecoder(filename string, decoder Decoder) (*CandidateNode, error) {
 	if decoder == nil {
 		return nil, fmt.Errorf("could not load %s", filename)
 	}
@@ -62,8 +61,39 @@ func loadYaml(filename string, decoder Decoder) (*CandidateNode, error) {
 	return sequenceNode, nil
 }
 
-func loadYamlOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
-	log.Debugf("loadYamlOperator")
+func loadStringOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
+	log.Debugf("loadString")
+
+	var results = list.New()
+
+	for el := context.MatchingNodes.Front(); el != nil; el = el.Next() {
+		candidate := el.Value.(*CandidateNode)
+
+		rhs, err := d.GetMatchingNodes(context.SingleReadonlyChildContext(candidate), expressionNode.RHS)
+		if err != nil {
+			return Context{}, err
+		}
+		if rhs.MatchingNodes.Front() == nil {
+			return Context{}, fmt.Errorf("filename expression returned nil")
+		}
+		nameCandidateNode := rhs.MatchingNodes.Front().Value.(*CandidateNode)
+
+		filename := nameCandidateNode.Value
+
+		contentsCandidate, err := loadString(filename)
+		if err != nil {
+			return Context{}, fmt.Errorf("Failed to load %v: %w", filename, err)
+		}
+
+		results.PushBack(contentsCandidate)
+
+	}
+
+	return context.ChildContext(results), nil
+}
+
+func loadOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
+	log.Debugf("loadOperator")
 
 	loadPrefs := expressionNode.Operation.Preferences.(loadPrefs)
 
@@ -86,13 +116,7 @@ func loadYamlOperator(d *dataTreeNavigator, context Context, expressionNode *Exp
 
 		filename := nameCandidateNode.Value
 
-		var contentsCandidate *CandidateNode
-
-		if loadPrefs.loadAsString {
-			contentsCandidate, err = loadString(filename)
-		} else {
-			contentsCandidate, err = loadYaml(filename, loadPrefs.decoder)
-		}
+		contentsCandidate, err := loadWithDecoder(filename, loadPrefs.decoder)
 		if err != nil {
 			return Context{}, fmt.Errorf("Failed to load %v: %w", filename, err)
 		}