From 061d8a8efc338305eb07cb57d1912b0fa7f121fa Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <pair+mbrunsfeld@pivotallabs.com>
Date: Wed, 6 Nov 2013 13:37:01 -0800
Subject: [PATCH] Compute transitions correctly for long sequences

---
 spec/rules_spec.cpp  | 10 ++++++++++
 src/rules.cpp        | 13 +++++++++----
 src/rules.h          | 10 ++++++----
 src/transition_map.h | 16 +++++++++++-----
 4 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/spec/rules_spec.cpp b/spec/rules_spec.cpp
index 98a9eb76..c878c060 100644
--- a/spec/rules_spec.cpp
+++ b/spec/rules_spec.cpp
@@ -12,6 +12,7 @@ Describe(Rules) {
         Symbol symbol1 = Symbol(1);
         Symbol symbol2 = Symbol(2);
         Symbol symbol3 = Symbol(3);
+        Symbol symbol4 = Symbol(3);
         
         It(handles_symbols) {
             AssertThat(
@@ -40,6 +41,15 @@ Describe(Rules) {
                 ), TransitionMap<Rule>::elements_equal));
         }
         
+        It(handles_long_sequences) {
+            AssertThat(
+                Seq(Seq(symbol1, symbol2), Seq(symbol3, symbol4)).transitions(),
+                EqualsContainer(TransitionMap<Rule>(
+                    { symbol1.copy() },
+                    { new Seq(symbol2, Seq(symbol3, symbol4)) }
+                ), TransitionMap<Rule>::elements_equal));
+        }
+
         It(handles_choices_with_common_starting_symbols) {
             AssertThat(
                 Choice(
diff --git a/src/rules.cpp b/src/rules.cpp
index 2f1408ac..8079678b 100644
--- a/src/rules.cpp
+++ b/src/rules.cpp
@@ -6,7 +6,6 @@ using namespace std;
 
 namespace tree_sitter  {
     namespace rules {
-        
         // Constructors
         Blank::Blank() {}
         Symbol::Symbol(int id) : id(id) {};
@@ -28,14 +27,20 @@ namespace tree_sitter  {
         
         TransitionMap<Rule> Choice::transitions() const {
             auto result = left->transitions();
-            result.merge(right->transitions(), [&](const Rule &left, const Rule &right) {
-                return new Choice(left, right);
+            result.merge(right->transitions(), [&](rule_ptr left, rule_ptr right) -> rule_ptr {
+                return rule_ptr(new Choice(left, right));
             });
             return result;
         }
         
         TransitionMap<Rule> Seq::transitions() const {
-            return TransitionMap<Rule>({ left->copy() }, { right->copy() });
+            return left->transitions().map([&](rule_ptr left_rule) -> rule_ptr {
+                if (typeid(*left_rule) == typeid(Blank)) {
+                    return right;
+                } else {
+                    return rule_ptr(new Seq(left_rule, right));
+                }
+            });
         }
         
         // Equality
diff --git a/src/rules.h b/src/rules.h
index 1e5e4a2e..f5f2a2e4 100644
--- a/src/rules.h
+++ b/src/rules.h
@@ -15,6 +15,8 @@ namespace tree_sitter {
             virtual std::string to_string() const = 0;
         };
         
+        typedef std::shared_ptr<const Rule> rule_ptr;
+        
         class Blank : public Rule {
         public:
             Blank();
@@ -45,8 +47,8 @@ namespace tree_sitter {
             bool operator==(const Rule& other) const;
             std::string to_string() const;
         private:
-            std::shared_ptr<const Rule> left;
-            std::shared_ptr<const Rule> right;
+            rule_ptr left;
+            rule_ptr right;
         };
         
         class Seq : public Rule {
@@ -59,8 +61,8 @@ namespace tree_sitter {
             bool operator==(const Rule& other) const;
             std::string to_string() const;
         private:
-            std::shared_ptr<const Rule> left;
-            std::shared_ptr<const Rule> right;
+            rule_ptr left;
+            rule_ptr right;
         };
     }
 }
diff --git a/src/transition_map.h b/src/transition_map.h
index 9a657312..f45b3ddf 100644
--- a/src/transition_map.h
+++ b/src/transition_map.h
@@ -66,16 +66,22 @@ namespace tree_sitter {
             return contents[i].second;
         }
         
-        void merge(const TransitionMap<MappedType> &other, std::function<const MappedType *(const MappedType &, const MappedType &)> merge_fn) {
+        void merge(const TransitionMap<MappedType> &other, std::function<mapped_ptr(mapped_ptr, mapped_ptr)> merge_fn) {
             for (pair_type other_pair : other) {
-                pair_type *current_pair = pair_for_key(*other_pair.first);
-                if (current_pair)
-                    current_pair->second = mapped_ptr(merge_fn(*current_pair->second, *other_pair.second));
+                if (pair_type *current_pair = pair_for_key(*other_pair.first))
+                    current_pair->second = merge_fn(current_pair->second, other_pair.second);
                 else
                     add(other_pair.first, other_pair.second);
             }
         }
-        
+
+        TransitionMap<MappedType> map(std::function<mapped_ptr(mapped_ptr)> map_fn) {
+            TransitionMap<MappedType> result;
+            for (pair_type pair : *this)
+                result.add(pair.first, map_fn(pair.second));
+            return result;
+        }
+
     private:
         pair_type * pair_for_key(rules::Rule const &key) {
             for (int i = 0; i < contents.size(); i++) {