Merge pull request #2507 from emmericp/visualize-cmd

Add --visualize command that outputs the AST of a given file as graphviz

Author: sumneko

script/cli/init.lua

@@ -12,3 +12,8 @@ if _G['DOC'] then
     require 'cli.doc' .runCLI()
     os.exit(0, true)
 end
+
+if _G['VISUALIZE'] then
+	local ret = require 'cli.visualize' .runCLI()
+	os.exit(ret or 0, true)
+end

script/cli/visualize.lua

@@ -0,0 +1,103 @@
+local lang   = require 'language'
+local parser = require 'parser'
+local guide  = require 'parser.guide'
+
+local function nodeId(node)
+	return node.type .. ':' .. node.start .. ':' .. node.finish
+end
+
+local function shorten(str)
+	if type(str) ~= 'string' then
+		return str
+	end
+	str = str:gsub('\n', '\\\\n')
+	if #str <= 20 then
+		return str
+	else
+		return str:sub(1, 17) .. '...'
+	end
+end
+
+local function getTooltipLine(k, v)
+	if type(v) == 'table' then
+		if v.type then
+			v = '<node ' .. v.type .. '>'
+		else
+			v = '<table>'
+		end
+	end
+	v = tostring(v)
+	v = v:gsub('"', '\\"')
+	return k .. ': ' .. shorten(v) .. '\\n'
+end
+
+local function getTooltip(node)
+	local str = ''
+	local skipNodes = {parent = true, start = true, finish = true, type = true}
+	str = str .. getTooltipLine('start', node.start)
+	str = str .. getTooltipLine('finish', node.finish)
+	for k, v in pairs(node) do
+		if type(k) ~= 'number' and not skipNodes[k] then
+			str = str .. getTooltipLine(k, v)
+		end
+	end
+	for i = 1, math.min(#node, 15) do
+		str = str .. getTooltipLine(i, node[i])
+	end
+	if #node > 15 then
+		str = str .. getTooltipLine('15..' .. #node, '(...)')
+	end
+	return str
+end
+
+local nodeEntry = '\t"%s" [\n\t\tlabel="%s\\l%s\\l"\n\t\ttooltip="%s"\n\t]'
+local function getNodeLabel(node)
+	local keyName = guide.getKeyName(node)
+	if node.type == 'binary' or node.type == 'unary' then
+		keyName = node.op.type
+	elseif node.type == 'label' or node.type == 'goto' then
+		keyName = node[1]
+	end
+	return nodeEntry:format(nodeId(node), node.type, shorten(keyName) or '', getTooltip(node))
+end
+
+local function getVisualizeVisitor(writer)
+	local function visitNode(node, parent)
+		if node == nil then return end
+		writer:write(getNodeLabel(node))
+		writer:write('\n')
+		if parent then
+			writer:write(('\t"%s" -> "%s"'):format(nodeId(parent), nodeId(node)))
+			writer:write('\n')
+		end
+		guide.eachChild(node, function(child)
+			visitNode(child, node)
+		end)
+	end
+	return visitNode
+end
+
+
+local export = {}
+
+function export.visualizeAst(code, writer)
+	local state = parser.compile(code, 'Lua', _G['LUA_VER'] or 'Lua 5.4')
+	writer:write('digraph AST {\n')
+	writer:write('\tnode [shape = rect]\n')
+	getVisualizeVisitor(writer)(state.ast)
+	writer:write('}\n')
+end
+
+function export.runCLI()
+	lang(LOCALE)
+	local file = _G['VISUALIZE']
+	local code, err = io.open(file)
+	if not code then
+		io.stderr:write('failed to open ' .. file .. ': ' .. err)
+		return 1
+	end
+	code = code:read('a')
+	return export.visualizeAst(code, io.stdout)
+end
+
+return export

test.lua

@@ -108,6 +108,7 @@ local function main()
     test 'tclient'
     test 'full'
     test 'plugins.test'
+	test 'cli.test'
 end
 
 loadAllLibs()

test/cli/test.lua

@@ -0,0 +1 @@
+require 'cli.visualize.test'

test/cli/visualize/test.lua

@@ -0,0 +1,23 @@
+local visualize = require 'cli.visualize'
+
+local testDataDir = 'test/cli/visualize/testdata/'
+
+local function TestVisualize(fileName)
+	local inputFile = testDataDir .. fileName .. '.txt'
+	local outputFile = testDataDir .. fileName .. '-expected.txt'
+	local output = ''
+	local writer = {}
+	function writer:write(text)
+		output = output .. text
+	end
+	visualize.visualizeAst(io.open(inputFile):read('a'), writer)
+	local expectedOutput = io.open(outputFile):read('a')
+	if expectedOutput ~= output then
+		-- uncomment this to update reference output
+		--io.open(outputFile, "w+"):write(output):close()
+		error('output mismatch for test file ' .. inputFile)
+	end
+end
+
+TestVisualize('all-types')
+TestVisualize('shorten-names')