fixes and performance optimization for transients

A bug was causing less in-place updates than possible

spec/immutable/vector/trie_spec.cr

@@ -258,7 +258,7 @@ describe Immutable::Vector::Trie do
           not_in_place += 1 unless x == t
           t = x
         end
-        not_in_place.should eq(Immutable::Vector::Trie::BLOCK_SIZE + 1)
+        not_in_place.should eq(2)
         t2.size.should eq(0)
         # update!
         t.update!(50, 0, 2_u64)

src/immutable/map/trie.cr

@@ -85,16 +85,16 @@ module Immutable
         self
       end
 
-      protected def set_at_index(index : Int32, key : K, value : V) : Trie(K, V)
+      protected def set_at_index(index : Int32, key : K, value : V, from = nil : UInt64?) : Trie(K, V)
         if leaf_of?(index)
-          set_leaf(index, key, value)
+          set_leaf(index, key, value, from)
         else
-          set_branch(index, key, value)
+          set_branch(index, key, value, from)
         end
       end
 
       protected def set_at_index!(index : Int32, key : K, value : V, from : UInt64) : Trie(K, V)
-        return set_at_index(index, key, value) unless from == @owner
+        return set_at_index(index, key, value, from) unless from == @owner
         if leaf_of?(index)
           @values[key] = value
         else
@@ -136,25 +136,25 @@ module Immutable
         (index.to_u32 >> (@levels * BITS_PER_LEVEL)) == 0
       end
 
-      private def set_leaf(index : Int32, key : K, value : V) : Trie(K, V)
+      private def set_leaf(index : Int32, key : K, value : V, from : UInt64?) : Trie(K, V)
         values = @values.dup.tap do |vs|
           vs[key] = value
         end
-        Trie.new(@children, values, @bitmap, @levels)
+        Trie.new(@children, values, @bitmap, @levels, from)
       end
 
-      private def set_branch(index : Int32, key : K, value : V) : Trie(K, V)
+      private def set_branch(index : Int32, key : K, value : V, from : UInt64?) : Trie(K, V)
         i = bit_index(index)
         if idx = child_index?(i)
           children = @children.dup.tap do |cs|
             cs[idx] = cs[idx].set_at_index(index, key, value)
           end
-          Trie.new(children, @values, @bitmap, @levels)
+          Trie.new(children, @values, @bitmap, @levels, from)
         else
-          child = Trie.new(([] of Trie(K, V)), Values(K, V).new, 0_u32, @levels + 1)
+          child = Trie.new(([] of Trie(K, V)), Values(K, V).new, 0_u32, @levels + 1, from)
             .set_at_index(index, key, value)
           bitmap = @bitmap | bitpos(i)
-          Trie.new(@children.dup.insert(child_index(i, bitmap), child), @values, bitmap, @levels)
+          Trie.new(@children.dup.insert(child_index(i, bitmap), child), @values, bitmap, @levels, from)
         end
       end
 

src/immutable/vector/trie.cr

@@ -108,25 +108,25 @@ module Immutable
         end.flatten
       end
 
-      def push_leaf(leaf : Array(T)) : Trie(T)
+      def push_leaf(leaf : Array(T), from = nil : UInt64?) : Trie(T)
         raise ArgumentError.new if leaf.size > BLOCK_SIZE || size % 32 != 0
-        return Trie.new([self], @levels + 1).push_leaf(leaf) if full?
-        return Trie.new(leaf) if empty? && leaf?
+        return Trie.new([self], @levels + 1, from).push_leaf(leaf, from) if full?
+        return Trie.new(leaf, from) if empty? && leaf?
         Trie.new(@children.dup.tap do |cs|
           if @levels == 1
             cs.push(Trie.new(leaf))
           else
             if cs.empty? || cs.last.full?
               cs << Trie.new([] of Trie(T), @levels - 1)
             end
-            cs[-1] = cs[-1].push_leaf(leaf)
+            cs[-1] = cs[-1].push_leaf(leaf, from)
           end
-        end, @levels)
+        end, @levels, from)
       end
 
       def push_leaf!(leaf : Array(T), from : UInt64) : Trie(T)
         raise ArgumentError.new if leaf.size > BLOCK_SIZE || size % 32 != 0
-        return push_leaf(leaf) unless from == @owner
+        return push_leaf(leaf, from) unless from == @owner
         return Trie.new([self], @levels + 1, from).push_leaf!(leaf, from) if full?
         return Trie.new(leaf, @owner) if empty? && leaf?
         if @levels == 1
@@ -141,20 +141,20 @@ module Immutable
         self
       end
 
-      def pop_leaf : Trie(T)
+      def pop_leaf(from = nil : UInt64?) : Trie(T)
         raise ArgumentError.new if empty? || size % 32 != 0
-        return Trie.new([] of T) if leaf?
+        return Trie.new([] of T, from) if leaf?
         child = @children.last.pop_leaf
         if child.empty?
           return @children.first if @children.size == 2
-          return Trie.new(@children[0...-1], @levels)
+          return Trie.new(@children[0...-1], @levels, from)
         end
-        Trie.new(@children[0...-1].push(child), @levels)
+        Trie.new(@children[0...-1].push(child), @levels, from)
       end
 
       def pop_leaf!(from : UInt64) : Trie(T)
         raise ArgumentError.new if empty? || size % 32 != 0
-        return pop_leaf unless from == @owner
+        return pop_leaf(from) unless from == @owner
         return Trie.new([] of T, from) if leaf?
         @children[-1] = @children[-1].pop_leaf!(from)
         if @children[-1].empty?
@@ -211,14 +211,14 @@ module Immutable
         trie.clear_owner!
       end
 
-      protected def set(index : Int, value : T) : Trie(T)
+      protected def set(index : Int, value : T, from = nil : UInt64?) : Trie(T)
         child_idx = child_index(index)
-        return Trie.new(update_values(child_idx, value)) if leaf?
-        Trie.new(update_children(child_idx, value, index), @levels)
+        return Trie.new(update_values(child_idx, value), from) if leaf?
+        Trie.new(update_children(child_idx, value, index), @levels, from)
       end
 
       protected def set!(index : Int, value : T, from : UInt64) : Trie(T)
-        return set(index, value) unless from == @owner
+        return set(index, value, from) unless from == @owner
         child_idx = child_index(index)
         if leaf?
           update_values!(child_idx, value, from)