X-Git-Url: http://nitlanguage.org

diff --git a/src/transform.nit b/src/transform.nit
index fe46105..2fda94e 100644
--- a/src/transform.nit
+++ b/src/transform.nit
@@ -68,8 +68,7 @@ private class TransformVisitor
 	redef fun visit(node)
 	do
 		if node isa AAnnotations then return
-		node.visit_all(self)
-		node.accept_transform_visitor(self)
+		node.full_transform_visitor(self)
 	end
 
 	# Get a primitive class or display a fatal error on `location`.
@@ -86,11 +85,32 @@ private class TransformVisitor
 end
 
 redef class ANode
+	private fun full_transform_visitor(v: TransformVisitor)
+	do
+		visit_all(v)
+		accept_transform_visitor(v)
+	end
 	private fun accept_transform_visitor(v: TransformVisitor)
 	do
 	end
 end
 
+redef class AExpr
+	redef fun full_transform_visitor(v: TransformVisitor)
+	do
+		var na = comprehension
+		if na != null then
+			# We are building a comprehension array `array`
+			# Replace `self` with `array.push(self)`
+			var place = detach_with_placeholder
+			var recv = na.nnew.make_var_read
+			var nadd = v.builder.make_call(recv, na.push_callsite.as(not null), [self])
+			place.replace_with(nadd)
+		end
+		super
+	end
+end
+
 redef class AVardeclExpr
 	# `var x = y` is replaced with `x = y`
 	#
@@ -278,22 +298,27 @@ redef class AArrayExpr
 	# t.add(y)
 	# t
 	# ~~~
-	redef fun accept_transform_visitor(v)
+	redef fun full_transform_visitor(v)
 	do
 		var nblock = v.builder.make_block
 
 		var nnew = v.builder.make_new(with_capacity_callsite.as(not null), [v.builder.make_int(n_exprs.length)])
+		self.nnew = nnew
+
 		nblock.add nnew
 
+		super
+
 		for nexpr in self.n_exprs do
-			var nadd = v.builder.make_call(nnew.make_var_read, push_callsite.as(not null), [nexpr])
-			nblock.add nadd
+			nblock.add nexpr
 		end
 		var nres = nnew.make_var_read
 		nblock.add nres
 
 		replace_with(nblock)
 	end
+
+	private var nnew: ANewExpr is noinit
 end
 
 redef class ACrangeExpr