lib/xmlss/element_stack.rb in xmlss-1.0.0.rc.1 vs lib/xmlss/element_stack.rb in xmlss-1.0.0.rc.2

@@ -10,11 +10,11 @@
     attr_reader :markup_type
 
     def initialize(writer, markup_type)
       @stack = []
       @writer = writer
-      @markup_type = markup_type
+      @markup_type = markup_type.to_s
       @written_level = 0
     end
 
     def empty?; @stack.empty?; end
     def size;   @stack.size;   end
@@ -22,47 +22,41 @@
     def last;   @stack.last;   end
 
     alias_method :current, :last
     alias_method :level, :size
 
+    def using(element, &block)
+      push(element)
+      (block || Proc.new {}).call
+      pop
+    end
+
     def push(element)
       if @written_level < level
-        open(current)
+        write(current)
+        @writer.push(@markup_type)
+        @written_level += 1
       end
       @stack.push(element)
     end
 
     def pop
       if !empty?
-        if @written_level < level
-          @stack.pop.tap { |elem| write(elem) }
-        else
-          @stack.pop.tap { |elem| close(elem) }
-        end
+        @written_level < level ? write(@stack.pop) : close(@stack.pop)
       end
     end
 
-    def using(element, &block)
-      push(element)
-      (block || Proc.new {}).call
-      pop
-    end
-
     private
 
-    def open(element)
-      write(element)
-      @writer.push(@markup_type)
-      @written_level += 1
-    end
-
     def close(element)
       @writer.pop(@markup_type)
       @written_level -= 1
+      element
     end
 
     def write(element)
       @writer.write(element)
+      element
     end
 
   end
 
 end