From f298eebb8b39d8bdefa57c33eec437c978f3bd1d Mon Sep 17 00:00:00 2001
From: zongz <68977949+zong-zhe@users.noreply.github.com>
Date: Mon, 21 Oct 2024 12:01:44 +0800
Subject: [PATCH] fix: fixed OverrideFile insert unification position error
 (#1702)

Signed-off-by: zongz <zongzhe1024@163.com>
---
 kclvm/query/src/override.rs        | 103 +++++++++++++++--------------
 kclvm/query/src/test_data/expect.k |   4 ++
 kclvm/query/src/tests.rs           |   2 +
 3 files changed, 59 insertions(+), 50 deletions(-)

diff --git a/kclvm/query/src/override.rs b/kclvm/query/src/override.rs
index 17456f0c2..871ee7efd 100644
--- a/kclvm/query/src/override.rs
+++ b/kclvm/query/src/override.rs
@@ -408,52 +408,66 @@ impl<'ctx> MutSelfMutWalker<'ctx> for OverrideTransformer {
                             }
                         }
                     } else if let ast::Stmt::Unification(unification_stmt) = &mut stmt.node {
-                        let target = match unification_stmt.target.node.names.get(0) {
-                            Some(name) => name,
-                            None => bug!(
-                                "Invalid AST unification target names {:?}",
-                                unification_stmt.target.node.names
-                            ),
-                        };
-                        if target.node == self.target_id {
-                            let item = unification_stmt.value.clone();
-                            let mut value = self.clone_override_value();
-                            // Use position information that needs to override the expression.
-                            value.set_pos(item.pos());
-                            let schema_expr = &mut unification_stmt.value.node;
-                            match &self.operation {
-                                ast::ConfigEntryOperation::Union => {
-                                    if let ast::Expr::Config(merged_config_expr) = &value.node {
-                                        if let ast::Expr::Config(config_expr) =
-                                            &mut schema_expr.config.node
-                                        {
-                                            self.has_override = merge_config_expr(
-                                                config_expr,
-                                                merged_config_expr,
-                                                &self.action,
-                                            );
-                                        }
-                                    } else if let ast::Expr::Schema(merged_schema_expr) =
-                                        &value.node
-                                    {
-                                        if schema_expr.name.node.get_name()
-                                            == merged_schema_expr.name.node.get_name()
-                                        {
-                                            if let (
-                                                ast::Expr::Config(merged_config_expr),
-                                                ast::Expr::Config(config_expr),
-                                            ) = (
-                                                &merged_schema_expr.config.node,
-                                                &mut schema_expr.config.node,
-                                            ) {
+                        if self.field_paths.len() == 0 {
+                            let target = match unification_stmt.target.node.names.get(0) {
+                                Some(name) => name,
+                                None => bug!(
+                                    "Invalid AST unification target names {:?}",
+                                    unification_stmt.target.node.names
+                                ),
+                            };
+                            if target.node == self.target_id {
+                                let item = unification_stmt.value.clone();
+                                let mut value = self.clone_override_value();
+                                // Use position information that needs to override the expression.
+                                value.set_pos(item.pos());
+                                let schema_expr = &mut unification_stmt.value.node;
+                                match &self.operation {
+                                    ast::ConfigEntryOperation::Union => {
+                                        if let ast::Expr::Config(merged_config_expr) = &value.node {
+                                            if let ast::Expr::Config(config_expr) =
+                                                &mut schema_expr.config.node
+                                            {
                                                 self.has_override = merge_config_expr(
                                                     config_expr,
                                                     merged_config_expr,
                                                     &self.action,
                                                 );
                                             }
+                                        } else if let ast::Expr::Schema(merged_schema_expr) =
+                                            &value.node
+                                        {
+                                            if schema_expr.name.node.get_name()
+                                                == merged_schema_expr.name.node.get_name()
+                                            {
+                                                if let (
+                                                    ast::Expr::Config(merged_config_expr),
+                                                    ast::Expr::Config(config_expr),
+                                                ) = (
+                                                    &merged_schema_expr.config.node,
+                                                    &mut schema_expr.config.node,
+                                                ) {
+                                                    self.has_override = merge_config_expr(
+                                                        config_expr,
+                                                        merged_config_expr,
+                                                        &self.action,
+                                                    );
+                                                }
+                                            }
+                                        } else {
+                                            // Unification is only support to override the schema expression.
+                                            if let ast::Expr::Schema(schema_expr) = value.node {
+                                                if self.field_paths.len() == 0 {
+                                                    unification_stmt.value = Box::new(
+                                                        ast::Node::dummy_node(schema_expr),
+                                                    );
+                                                    self.has_override = true;
+                                                }
+                                            }
                                         }
-                                    } else {
+                                    }
+                                    ast::ConfigEntryOperation::Insert
+                                    | ast::ConfigEntryOperation::Override => {
                                         // Unification is only support to override the schema expression.
                                         if let ast::Expr::Schema(schema_expr) = value.node {
                                             if self.field_paths.len() == 0 {
@@ -464,17 +478,6 @@ impl<'ctx> MutSelfMutWalker<'ctx> for OverrideTransformer {
                                         }
                                     }
                                 }
-                                ast::ConfigEntryOperation::Insert
-                                | ast::ConfigEntryOperation::Override => {
-                                    // Unification is only support to override the schema expression.
-                                    if let ast::Expr::Schema(schema_expr) = value.node {
-                                        if self.field_paths.len() == 0 {
-                                            unification_stmt.value =
-                                                Box::new(ast::Node::dummy_node(schema_expr));
-                                            self.has_override = true;
-                                        }
-                                    }
-                                }
                             }
                         }
                     }
diff --git a/kclvm/query/src/test_data/expect.k b/kclvm/query/src/test_data/expect.k
index 5b3e730b2..90ca23838 100644
--- a/kclvm/query/src/test_data/expect.k
+++ b/kclvm/query/src/test_data/expect.k
@@ -61,4 +61,8 @@ uni_config = {
 
 config_unification: Config {
     "image": "image/image:v4"
+    env: {
+        "aaa": "aaa"
+        bbb: "bbb"
+    }
 }
diff --git a/kclvm/query/src/tests.rs b/kclvm/query/src/tests.rs
index f18de5142..77c54fdbd 100644
--- a/kclvm/query/src/tests.rs
+++ b/kclvm/query/src/tests.rs
@@ -43,6 +43,8 @@ fn test_override_file_simple() {
         "insert_config.key=1".to_string(),
         "uni_config.labels.key1=1".to_string(),
         "config_unification=Config {\"image\": \"image/image:v4\"}".to_string(),
+        "config_unification:Config {\"env\": {\"aaa\": \"aaa\"}}".to_string(),
+        "config_unification.env: {\"bbb\": \"bbb\"}}".to_string(),
         "config_unification_delete-".to_string()
     ];