Replace AttachBlock with more generic PartialScope

ast/AttachBlock.lua

@@ -1,64 +0,0 @@
-local ast = require("ast")
-local Identifier, Quote
-
-local attached_block_identifier, attached_block_symbol
-
-local AttachBlock
-AttachBlock = ast.abstract.Node {
-	type = "attach block",
-
-	expression = nil,
-	block = nil,
-
-	init = function(self, expression, block)
-		self.expression = expression
-		self.block = block
-		self.format_priority = self.expression.format_priority
-	end,
-
-	_format = function(self, state, priority, indentation, ...)
-		local exp = self.expression:format(state, priority, indentation, ...)
-		if exp:sub(-2) == " _" then exp = exp:sub(1, -3) end
-		return exp.."\n\t"..self.block:format(state, priority, indentation + 1, ...)
-	end,
-
-	traverse = function(self, fn, ...)
-		fn(self.expression, ...)
-		fn(self.block, ...)
-	end,
-
-	_eval = function(self, state)
-		state.scope:push_partial(attached_block_identifier)
-		state.scope:define(attached_block_symbol, Quote:new(self.block)) -- _ is always wrapped in a Call when it appears
-		local exp = self.expression:eval(state)
-		state.scope:pop()
-
-		return exp
-	end,
-
-	_prepare = function(self, state)
-		state.scope:push_partial(attached_block_identifier)
-		state.scope:define(attached_block_symbol, Quote:new(self.block))
-		self.expression:prepare(state)
-		state.scope:pop()
-	end,
-
-	-- class method: if the block identifier is defined in the current scope, wrap node in an AttachBlock so the block is still defined in this node
-	-- used to preserve the defined _ block without the need to build a full closure
-	-- used e.g. for -> translation, as we want to preserve _ while still executing the translation in the Translatable scope and not restore a different scope from a closure
-	-- (operates on un-evaluated nodes!)
-	preserve = function(self, state, node)
-		if state.scope:defined_in_current(attached_block_symbol) then
-			return AttachBlock:new(node, state.scope:get(attached_block_identifier).expression) -- unwrap Quote as that will be rewrap on eval
-		end
-		return node
-	end,
-}
-
-package.loaded[...] = AttachBlock
-Identifier, Quote = ast.Identifier, ast.Quote
-
-attached_block_identifier = Identifier:new("_")
-attached_block_symbol = attached_block_identifier:to_symbol()
-
-return AttachBlock

ast/Environment.lua

@@ -22,6 +22,9 @@ local VariableMetadata = ast.abstract.Runtime {
 			return self.branched:get(state)
 		end
 	end,
+	get_symbol = function(self)
+		return self.symbol
+	end,
 	set = function(self, state, value)
 		if self.symbol.constant then
 			error(("trying to change the value of constant %s"):format(self.symbol.string), 0)
@@ -180,6 +183,10 @@ local Environment = ast.abstract.Runtime {
 	get = function(self, state, identifier)
 		return self:_get_variable(state, identifier):get(state)
 	end,
+	-- get the symbol that was used to define the variable in current or parent environment
+	get_symbol = function(self, state, identifier)
+		return self:_get_variable(state, identifier):get_symbol()
+	end,
 	-- set variable value in current or parent environment
 	set = function(self, state, identifier, val)
 		return self:_get_variable(state, identifier):set(state, val)

ast/Function.lua

@@ -33,6 +33,10 @@ Function = Overloadable {
 	traverse = function(self, fn, ...)
 		fn(self.parameters, ...)
 		fn(self.expression, ...)
+		for sym, val in pairs(self.exports) do
+			fn(sym, ...)
+			fn(val, ...)
+		end
 	end,
 
 	compatible_with_arguments = function(self, state, args)

ast/PartialScope.lua

@@ -0,0 +1,90 @@
+-- create a partial layer to define temporary variables
+
+local ast = require("ast")
+local Identifier, Quote
+
+local attached_block_identifier, attached_block_symbol
+
+local PartialScope
+PartialScope = ast.abstract.Node {
+	type = "partial scope",
+
+	expression = nil,
+	definitions = nil, -- {[sym]=value,...} where values are already evaluated!
+	_identifiers = nil, -- {identifier,...} - just a cache so we don't rebuild it on every eval
+
+	init = function(self, expression)
+		self.expression = expression
+		self.definitions = {}
+		self._identifiers = {}
+		self.format_priority = self.expression.format_priority
+	end,
+	define = function(self, symbol, value) -- for construction only
+		assert(not self.definitions[symbol], ("%s already defined in partial layer"):format(symbol))
+		table.insert(self._identifiers, symbol:to_identifier())
+		self.definitions[symbol] = value
+	end,
+
+	_format = function(self, state, priority, indentation, ...)
+		if self.definitions[attached_block_symbol] then
+			local block = self.definitions[attached_block_symbol]
+			local exp = self.expression:format(state, priority, indentation, ...)
+			if exp:sub(-2) == " _" then exp = exp:sub(1, -3) end
+			return exp.."\n\t"..block:format(state, priority, indentation + 1, ...)
+		else
+			return self.expression:format(state, priority, indentation, ...)
+		end
+	end,
+
+	traverse = function(self, fn, ...)
+		fn(self.expression, ...)
+		for sym, val in pairs(self.definitions) do
+			fn(sym, ...)
+			fn(val, ...)
+		end
+	end,
+
+	_eval = function(self, state)
+		state.scope:push_partial(table.unpack(self._identifiers))
+		for sym, val in pairs(self.definitions) do state.scope:define(sym, val) end
+		local exp = self.expression:eval(state)
+		state.scope:pop()
+
+		return exp
+	end,
+
+	_prepare = function(self, state)
+		state.scope:push_partial(table.unpack(self._identifiers))
+		for sym, val in pairs(self.definitions) do state.scope:define(sym, val) end
+		self.expression:prepare(state)
+		state.scope:pop()
+	end,
+
+	-- class method: if the identifier is currently defined, wrap node in an PartialScope so the identifier is still defined in this node
+	-- used to e.g. preserve the defined _ block without the need to build a full closure
+	-- used e.g. for -> translation, as we want to preserve _ while still executing the translation in the Translatable scope and not restore a different scope from a closure
+	-- (operates on un-evaluated nodes!)
+	preserve = function(self, state, expression, ...)
+		local partial = PartialScope:new(expression)
+		for _, ident in ipairs{...} do
+			if state.scope:defined(ident) then
+				partial:define(state.scope:get_symbol(ident), state.scope:get(ident))
+			end
+		end
+		return partial
+	end,
+	-- class method: return a PartialScope that define the block identifier _ to a Quote of `block`
+	attach_block = function(self, expression, block)
+		local partial = ast.PartialScope:new(expression)
+		partial:define(attached_block_symbol, Quote:new(block))
+		return partial
+	end
+}
+
+package.loaded[...] = PartialScope
+Identifier, Quote = ast.Identifier, ast.Quote
+
+attached_block_identifier = Identifier:new("_")
+attached_block_symbol = attached_block_identifier:to_symbol()
+
+return PartialScope