Allow field names as rule options

Change-Id: Ife4f15a09818cf6daf86b96e7c916854d0299be8
diff --git a/mapper/corpus.go b/mapper/corpus.go
index f5800ba..9aad21e 100644
--- a/mapper/corpus.go
+++ b/mapper/corpus.go
@@ -10,7 +10,7 @@
 // Rules are applied iteratively: each rule is applied to the entire tree,
 // and subsequent rules see the transformed result.
 func (m *Mapper) applyCorpusQueryMappings(mappingID string, opts MappingOptions, jsonData any) (any, error) {
-	rules := m.parsedCorpusRules[mappingID]
+	rules := m.rulesWithFieldOverrides(m.parsedCorpusRules[mappingID], opts)
 
 	jsonMap, ok := jsonData.(map[string]any)
 	if !ok {
@@ -383,7 +383,7 @@
 
 // applyCorpusResponseMappings processes fields arrays with corpus rules.
 func (m *Mapper) applyCorpusResponseMappings(mappingID string, opts MappingOptions, jsonData any) (any, error) {
-	rules := m.parsedCorpusRules[mappingID]
+	rules := m.rulesWithFieldOverrides(m.parsedCorpusRules[mappingID], opts)
 
 	jsonMap, ok := jsonData.(map[string]any)
 	if !ok {
@@ -535,3 +535,40 @@
 	return result
 }
 
+func (m *Mapper) rulesWithFieldOverrides(rules []*parser.CorpusMappingResult, opts MappingOptions) []*parser.CorpusMappingResult {
+	if opts.FieldA == "" && opts.FieldB == "" {
+		return rules
+	}
+
+	result := make([]*parser.CorpusMappingResult, len(rules))
+	for i, rule := range rules {
+		upper := rule.Upper.Clone()
+		lower := rule.Lower.Clone()
+
+		if opts.FieldA != "" {
+			applyCorpusKeyOverride(upper, opts.FieldA)
+		}
+		if opts.FieldB != "" {
+			applyCorpusKeyOverride(lower, opts.FieldB)
+		}
+
+		result[i] = &parser.CorpusMappingResult{
+			Upper: upper,
+			Lower: lower,
+		}
+	}
+
+	return result
+}
+
+func applyCorpusKeyOverride(node parser.CorpusNode, key string) {
+	switch n := node.(type) {
+	case *parser.CorpusField:
+		n.Key = key
+	case *parser.CorpusGroup:
+		for _, op := range n.Operands {
+			applyCorpusKeyOverride(op, key)
+		}
+	}
+}
+
diff --git a/mapper/corpus_test.go b/mapper/corpus_test.go
index df2ae85..c5e129c 100644
--- a/mapper/corpus_test.go
+++ b/mapper/corpus_test.go
@@ -162,6 +162,66 @@
 	assert.Equal(t, "fiction", corpus["value"])
 }
 
+func TestCorpusQueryFieldOverridesAtoB(t *testing.T) {
+	m, err := NewMapper([]config.MappingList{{
+		ID:       "corpus-test",
+		Type:     "corpus",
+		FieldA:   "textClass",
+		FieldB:   "genre",
+		Mappings: []config.MappingRule{"novel <> fiction"},
+	}})
+	require.NoError(t, err)
+
+	input := map[string]any{
+		"corpus": map[string]any{
+			"@type": "koral:doc",
+			"key":   "domain",
+			"value": "novel",
+			"match": "match:eq",
+		},
+	}
+	result, err := m.ApplyQueryMappings("corpus-test", MappingOptions{
+		Direction: AtoB,
+		FieldA:    "domain",
+		FieldB:    "subject",
+	}, input)
+	require.NoError(t, err)
+
+	corpus := result.(map[string]any)["corpus"].(map[string]any)
+	assert.Equal(t, "subject", corpus["key"])
+	assert.Equal(t, "fiction", corpus["value"])
+}
+
+func TestCorpusQueryFieldOverridesBtoA(t *testing.T) {
+	m, err := NewMapper([]config.MappingList{{
+		ID:       "corpus-test",
+		Type:     "corpus",
+		FieldA:   "textClass",
+		FieldB:   "genre",
+		Mappings: []config.MappingRule{"novel <> fiction"},
+	}})
+	require.NoError(t, err)
+
+	input := map[string]any{
+		"corpus": map[string]any{
+			"@type": "koral:doc",
+			"key":   "subject",
+			"value": "fiction",
+			"match": "match:eq",
+		},
+	}
+	result, err := m.ApplyQueryMappings("corpus-test", MappingOptions{
+		Direction: BtoA,
+		FieldA:    "domain",
+		FieldB:    "subject",
+	}, input)
+	require.NoError(t, err)
+
+	corpus := result.(map[string]any)["corpus"].(map[string]any)
+	assert.Equal(t, "domain", corpus["key"])
+	assert.Equal(t, "novel", corpus["value"])
+}
+
 func TestCorpusQueryFieldGroupAlias(t *testing.T) {
 	m := newCorpusMapper(t, "textClass=novel <> genre=fiction")
 
diff --git a/mapper/mapper.go b/mapper/mapper.go
index 4869c7b..ce172de 100644
--- a/mapper/mapper.go
+++ b/mapper/mapper.go
@@ -100,6 +100,8 @@
 	LayerA      string
 	FoundryB    string
 	LayerB      string
+	FieldA      string
+	FieldB      string
 	Direction   Direction
 	AddRewrites bool
 }