Add utils for token modification in fyaml

add utils for token modification in fyaml

3 files changed, 422 insertions, 20 deletions

diff --git a/library/cpp/yaml/fyamlcpp/fyamlcpp.cpp b/library/cpp/yaml/fyamlcpp/fyamlcpp.cpp
index dc6caaf726..64e423f55d 100644
--- a/library/cpp/yaml/fyamlcpp/fyamlcpp.cpp
+++ b/library/cpp/yaml/fyamlcpp/fyamlcpp.cpp
@@ -318,9 +318,8 @@ TString TNodeRef::Scalar() const {
     return TString(text, size);
 }
 
-TMark TNodeRef::ScalarBeginMark() const {
+TMark TNodeRef::BeginMark() const {
     ENSURE_NODE_NOT_EMPTY(Node_);
-    Y_ENSURE_EX(fy_node_is_scalar(Node_), TFyamlEx() << "Node is not Scalar: " << Path());
     std::unique_ptr<fy_document_iterator, void(*)(fy_document_iterator*)> it(
         fy_document_iterator_create(),
         &fy_document_iterator_destroy);
@@ -329,8 +328,7 @@ TMark TNodeRef::ScalarBeginMark() const {
     std::unique_ptr<fy_event, decltype(deleter)> ev(
         fy_document_iterator_body_next(it.get()),
         deleter);
-    auto* token = fy_event_get_token(ev.get());
-    auto* mark = fy_token_start_mark(token);
+    auto* mark = fy_event_start_mark(ev.get());
     return TMark{
         mark->input_pos,
         mark->line,
@@ -338,19 +336,73 @@ TMark TNodeRef::ScalarBeginMark() const {
     };
 }
 
-TMark TNodeRef::ScalarEndMark() const {
+bool IsComplexType(ENodeType type) {
+    return type == ENodeType::Mapping || type == ENodeType::Sequence;
+}
+
+fy_event_type GetOpenEventType(ENodeType type) {
+    switch(type) {
+    case ENodeType::Mapping:
+        return FYET_MAPPING_START;
+    case ENodeType::Sequence:
+        return FYET_SEQUENCE_START;
+    default:
+        Y_FAIL("Not a brackets type");
+    }
+}
+
+fy_event_type GetCloseEventType(ENodeType type) {
+    switch(type) {
+    case ENodeType::Mapping:
+        return FYET_MAPPING_END;
+    case ENodeType::Sequence:
+        return FYET_SEQUENCE_END;
+    default:
+        Y_FAIL("Not a brackets type");
+    }
+}
+
+TMark TNodeRef::EndMark() const {
     ENSURE_NODE_NOT_EMPTY(Node_);
-    Y_ENSURE_EX(fy_node_is_scalar(Node_), TFyamlEx() << "Node is not Scalar: " << Path());
     std::unique_ptr<fy_document_iterator, void(*)(fy_document_iterator*)> it(
         fy_document_iterator_create(),
         &fy_document_iterator_destroy);
     fy_document_iterator_node_start(it.get(), Node_);
+
     auto deleter = [&](fy_event* fye){ fy_document_iterator_event_free(it.get(), fye); };
     std::unique_ptr<fy_event, decltype(deleter)> ev(
         fy_document_iterator_body_next(it.get()),
         deleter);
-    auto* token = fy_event_get_token(ev.get());
-    auto* mark = fy_token_end_mark(token);
+
+    if (IsComplexType(Type())) {
+        int openBrackets = 0;
+        if (ev->type == GetOpenEventType(Type())) {
+            ++openBrackets;
+        }
+        if (ev->type == GetCloseEventType(Type())) {
+            --openBrackets;
+        }
+        while (ev->type != GetCloseEventType(Type()) || openBrackets != 0) {
+            std::unique_ptr<fy_event, decltype(deleter)> cur(
+                fy_document_iterator_body_next(it.get()),
+                deleter);
+            if (cur == nullptr) {
+                break;
+            }
+            if (cur->type == GetOpenEventType(Type())) {
+                ++openBrackets;
+            }
+            if (cur->type == GetCloseEventType(Type())) {
+                --openBrackets;
+            }
+            if (fy_event_get_node_style(cur.get()) != FYNS_BLOCK) {
+                ev.reset(cur.release());
+            }
+        }
+    }
+
+    const auto* mark = fy_event_end_mark(ev.get());
+
     return TMark{
         mark->input_pos,
         mark->line,
@@ -467,6 +519,12 @@ TNodeRef TNodePairRef::Key() const {
     return TNodeRef(fy_node_pair_key(Pair_));
 }
 
+int TNodePairRef::Index(const TNodeRef& node) const {
+    ENSURE_NODE_NOT_EMPTY(node);
+    ENSURE_NODE_NOT_EMPTY(Pair_);
+    return fy_node_mapping_get_pair_index(node.Node_, Pair_);
+}
+
 void TNodePairRef::SetKey(const TNodeRef& node) {
     ENSURE_NODE_NOT_EMPTY(Pair_);
     ENSURE_NODE_NOT_EMPTY(node);
diff --git a/library/cpp/yaml/fyamlcpp/fyamlcpp.h b/library/cpp/yaml/fyamlcpp/fyamlcpp.h
index 7044a010a7..6df7b244c6 100644
--- a/library/cpp/yaml/fyamlcpp/fyamlcpp.h
+++ b/library/cpp/yaml/fyamlcpp/fyamlcpp.h
@@ -150,9 +150,9 @@ public:
 
     TString Scalar() const;
 
-    TMark ScalarBeginMark() const;
+    TMark BeginMark() const;
 
-    TMark ScalarEndMark() const;
+    TMark EndMark() const;
 
     void Insert(const TNodeRef& node);
 
@@ -219,6 +219,8 @@ public:
 
     TNodeRef Key() const;
 
+    int Index(const TNodeRef& node) const;
+
     void SetKey(const TNodeRef& node);
 
     TNodeRef Value() const;
diff --git a/library/cpp/yaml/fyamlcpp/fyamlcpp_ut.cpp b/library/cpp/yaml/fyamlcpp/fyamlcpp_ut.cpp
index 299fe614a2..2f6e14138c 100644
--- a/library/cpp/yaml/fyamlcpp/fyamlcpp_ut.cpp
+++ b/library/cpp/yaml/fyamlcpp/fyamlcpp_ut.cpp
@@ -164,22 +164,364 @@ x: b
         UNIT_ASSERT_VALUES_EQUAL(seq[1].Map().at("b").Sequence().at(2).Scalar(), "3");
     }
 
-    Y_UNIT_TEST(ScalarMark) {
-        TString str = R"(
+    Y_UNIT_TEST(SimpleScalarMark) {
+        auto check = [](const TString& str, const NFyaml::TNodeRef& node) {
+            auto pos = str.find("123");
+            auto endPos = pos + strlen("123");
+            auto begin = node.BeginMark().InputPos;
+            auto end = node.EndMark().InputPos;
+            UNIT_ASSERT_VALUES_EQUAL(begin, pos);
+            UNIT_ASSERT_VALUES_EQUAL(end, endPos);
+        };
+
+        {
+            TString str = R"(obj: 123)";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("obj");
+            check(str, node);
+        }
+
+        {
+            TString str = R"(obj: 123 # test)";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("obj");
+            check(str, node);
+        }
+
+        {
+            TString str = R"(
+# test
+obj: 123 # test
+# test
+)";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("obj");
+            check(str, node);
+        }
+
+        {
+            TString str = R"(
+---
+obj: 123
+...
+)";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("obj");
+            check(str, node);
+        }
+
+        {
+            TString str = R"(
 a: foo
 test: [{obj: 123}]
 b: bar
-        )";
+            )";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("test").Sequence().at(0).Map().at("obj");
+            check(str, node);
+        }
 
-        auto doc = NFyaml::TDocument::Parse(str);
+        {
+            TString str = R"(obj: '123')";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("obj");
+            check(str, node);
+        }
 
-        auto node = doc.Root().Map().at("test").Sequence().at(0).Map().at("obj");
+        {
+            TString str = R"(obj: '123' # test)";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("obj");
+            check(str, node);
+        }
 
-        auto pos = str.find("123");
-        auto begin = node.ScalarBeginMark().InputPos;
-        auto end = node.ScalarEndMark().InputPos;
+        {
+            TString str = R"(
+# test
+obj: '123' # test
+# test
+)";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("obj");
+            check(str, node);
+        }
 
-        UNIT_ASSERT_VALUES_EQUAL(begin, pos);
-        UNIT_ASSERT_VALUES_EQUAL(end, pos + strlen("123"));
+        {
+            TString str = R"(
+---
+obj: '123'
+...
+)";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("obj");
+            check(str, node);
+        }
+
+        {
+            TString str = R"(
+a: foo
+test: [{obj: "123"}]
+b: bar
+            )";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("test").Sequence().at(0).Map().at("obj");
+            check(str, node);
+        }
+
+        {
+            TString str = R"(obj: "123")";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("obj");
+            check(str, node);
+        }
+
+        {
+            TString str = R"(obj: "123" # test)";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("obj");
+            check(str, node);
+        }
+
+        {
+            TString str = R"(
+# test
+obj: "123" # test
+# test
+)";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("obj");
+            check(str, node);
+        }
+
+        {
+            TString str = R"(
+---
+obj: "123"
+...
+)";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("obj");
+            check(str, node);
+        }
+
+        {
+            TString str = R"(
+a: foo
+test: [{obj: "123"}]
+b: bar
+            )";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("test").Sequence().at(0).Map().at("obj");
+            check(str, node);
+        }
+
+        {
+            TString str = R"(
+a: foo
+test: [{obj: !!int "123"}]
+b: bar
+            )";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("test").Sequence().at(0).Map().at("obj");
+            check(str, node);
+        }
+
+    }
+
+    Y_UNIT_TEST(MultilineScalarMark) {
+        {
+            TString str = R"(obj: >+2
+  some
+  multiline
+
+  scalar with couple words)";
+            auto doc = NFyaml::TDocument::Parse(str);
+            auto node = doc.Root().Map().at("obj");
+            auto begin = node.BeginMark().InputPos;
+            auto end = node.EndMark().InputPos;
+            UNIT_ASSERT_VALUES_EQUAL(begin, 9);
+            UNIT_ASSERT_VALUES_EQUAL(end, 55);
+        }
     }
+
+    Y_UNIT_TEST(MapMark) {
+        {
+            TString str = R"(
+a: foo
+map: !!map
+  internal_map1: {}
+  internal_map2:
+    internal_map3:
+      internal_map4:
+        value: 1
+      internal_map5: {
+        internal_map6: {test1: 1, test2: 2},
+        internal_map7: {
+          value: 1
+          }
+        }
+# comment
+c: bar
+            )";
+
+            auto doc = NFyaml::TDocument::Parse(str);
+
+            auto node = doc.Root().Map().at("map");
+
+            auto begin = node.BeginMark().InputPos;
+            auto end = node.EndMark().InputPos;
+
+            UNIT_ASSERT_VALUES_EQUAL(begin, 21);
+            UNIT_ASSERT_VALUES_EQUAL(end, 246);
+        }
+
+        {
+            TString str = R"(
+a: foo
+map: !!map # comment
+# comment
+c: bar
+            )";
+
+            auto doc = NFyaml::TDocument::Parse(str);
+
+            auto node = doc.Root().Map().at("map");
+
+            auto begin = node.BeginMark().InputPos;
+            auto end = node.EndMark().InputPos;
+
+            UNIT_ASSERT_VALUES_EQUAL(begin, 11);
+            UNIT_ASSERT_VALUES_EQUAL(end, 11);
+        }
+
+        {
+            TString str = R"(
+a: foo
+map: {} # comment
+# comment
+c: bar
+            )";
+
+            auto doc = NFyaml::TDocument::Parse(str);
+
+            auto node = doc.Root().Map().at("map");
+
+            auto begin = node.BeginMark().InputPos;
+            auto end = node.EndMark().InputPos;
+
+            UNIT_ASSERT_VALUES_EQUAL(begin, 13);
+            UNIT_ASSERT_VALUES_EQUAL(end, 15);
+        }
+
+        {
+            TString str = R"(
+a: foo
+map:
+  value: 1
+# comment
+c: bar
+            )";
+
+            auto doc = NFyaml::TDocument::Parse(str);
+
+            auto node = doc.Root().Map().at("map");
+
+            auto begin = node.BeginMark().InputPos;
+            auto end = node.EndMark().InputPos;
+
+            UNIT_ASSERT_VALUES_EQUAL(begin, 15);
+            UNIT_ASSERT_VALUES_EQUAL(end, 23);
+        }
+    }
+
+    Y_UNIT_TEST(SequenceMark) {
+        {
+            TString str = R"(
+a: foo
+seq: !!seq
+- internal_map1: {}
+- internal_seq2:
+  - internal_seq3:
+    - internal_seq4:
+        value: 1
+    - internal_seq5: [
+        internal_seq6: [{test1: 1}, {test2: 2}],
+        internal_seq7: [
+          {value: 1}
+          ]
+        ]
+# comment
+c: bar
+            )";
+
+            auto doc = NFyaml::TDocument::Parse(str);
+
+            auto node = doc.Root().Map().at("seq");
+
+            auto begin = node.BeginMark().InputPos;
+            auto end = node.EndMark().InputPos;
+
+            UNIT_ASSERT_VALUES_EQUAL(begin, 19);
+            UNIT_ASSERT_VALUES_EQUAL(end, 252);
+        }
+
+        {
+            TString str = R"(
+a: foo
+seq: !!seq # comment
+# comment
+c: bar
+            )";
+
+            auto doc = NFyaml::TDocument::Parse(str);
+
+            auto node = doc.Root().Map().at("seq");
+
+            auto begin = node.BeginMark().InputPos;
+            auto end = node.EndMark().InputPos;
+
+            UNIT_ASSERT_VALUES_EQUAL(begin, 11);
+            UNIT_ASSERT_VALUES_EQUAL(end, 11);
+        }
+
+        {
+            TString str = R"(
+a: foo
+seq: [] # comment
+# comment
+c: bar
+            )";
+
+            auto doc = NFyaml::TDocument::Parse(str);
+
+            auto node = doc.Root().Map().at("seq");
+
+            auto begin = node.BeginMark().InputPos;
+            auto end = node.EndMark().InputPos;
+
+            UNIT_ASSERT_VALUES_EQUAL(begin, 13);
+            UNIT_ASSERT_VALUES_EQUAL(end, 15);
+        }
+
+        {
+            TString str = R"(
+a: foo
+seq:
+- value: 1
+# comment
+c: bar
+            )";
+
+            auto doc = NFyaml::TDocument::Parse(str);
+
+            auto node = doc.Root().Map().at("seq");
+
+            auto begin = node.BeginMark().InputPos;
+            auto end = node.EndMark().InputPos;
+
+            UNIT_ASSERT_VALUES_EQUAL(begin, 13);
+            UNIT_ASSERT_VALUES_EQUAL(end, 23);
+        }
+    }
+
 }