From e0552daf91e27776d2a8fa864e1604309222fab5 Mon Sep 17 00:00:00 2001
From: Jon Michael Aanes <jonjmaa@gmail.com>
Date: Sat, 22 Jul 2017 15:09:20 +0200
Subject: [PATCH] Chunks are now formatted correctly.

---
 function.lua           | 41 +++++++++++++++++++++++++++--------------
 test/test_function.lua | 14 ++++++++++++++
 test/test_pretty.lua   |  6 ------
 3 files changed, 41 insertions(+), 20 deletions(-)

diff --git a/function.lua b/function.lua
index a052c83..c42f2fb 100644
--- a/function.lua
+++ b/function.lua
@@ -195,27 +195,40 @@ local function get_function_body_info (info)
 
 	if not str then  return info  end
 
-	-- Calculate indices of the lines the function should be defined at.
-    local start_line_index  =  get_line_index(str, info.linedefined)
-    local end_line_index    =  get_line_index(str, info.lastlinedefined + 1)
-
 	-- Now find some info about the function.
 	-- NOTE: function_params is currently not used for anything.
-	local function_name, function_params, function_body  =  str:sub(start_line_index, end_line_index):match(FUNCTION_DEFINITION_MATCH)
+	local function_name, function_params, function_body
 
-	if type(function_body) ~= 'string' then
-		error(('[pretty.function/internal]: Could not find the function defined on lines %i-%i (indices %i-%i) for string:\n\n%s\n'):format(info.linedefined, info.lastlinedefined, start_line_index, end_line_index, str))
+	if info.linedefined == 0 then
+		-- A function is a "chunk" when linedefined is equal to 0.
+		-- This is a "toplevel" function. One without a function
+		-- definition. The entire string is it's function body.
+		function_name	= ''
+		function_body	= str
+	else
+		-- This is not a chunk. Look for function definition.
+
+		-- Calculate indices of the lines the function should be defined at.
+		local start_line_index	=  get_line_index(str, info.linedefined)
+		local end_line_index	=  get_line_index(str, info.lastlinedefined + 1)
+
+		function_name, function_params, function_body  =  str:sub(start_line_index, end_line_index):match(FUNCTION_DEFINITION_MATCH)
+
+		-- Throw an error if we can't find anything.
+		if type(function_body) ~= 'string' then
+			error(('[pretty.function/internal]: Could not find the function defined on lines %i-%i (indices %i-%i) for string:\n\n%s\n'):format(info.linedefined, info.lastlinedefined, start_line_index, end_line_index, str))
+		end
 	end
 
-	local function_body  =  function_body:match '^%s*(.-)%s*$'
-	local pivot_index    =  get_docs_split_index (function_body)
+	local function_body	=  function_body:match '^%s*(.-)%s*$'
+	local pivot_index	=  get_docs_split_index (function_body)
 
-	info.name        =  info.name or function_name:match '^%s*(.-)%s*$'
-	info.docs         = info.docs or get_docs_from_function_body(function_body, pivot_index)
-	info.body         = info.body or function_body:sub(pivot_index, -1):match '^%s*(.-)%s*$'
+	info.name	= info.name or function_name:match '^%s*(.-)%s*$'
+	info.docs	= info.docs or get_docs_from_function_body(function_body, pivot_index)
+	info.body	= info.body or function_body:sub(pivot_index, -1):match '^%s*(.-)%s*$'
 
-	if info.name == '' then  info.name = nil  end
-	if info.docs == '' then  info.docs = nil  end
+	if info.name == '' then	info.name = nil	end
+	if info.docs == '' then	info.docs = nil	end
 
 	return info
 end
diff --git a/test/test_function.lua b/test/test_function.lua
index faa257d..f844693 100644
--- a/test/test_function.lua
+++ b/test/test_function.lua
@@ -269,6 +269,20 @@ format_test {
     expect = 'function (a, b) return a + b end',
 }
 
+format_test {
+    name   = 'Embedding loaded chunk',
+    single = true,
+    input  = loadstring 'return 42',
+    expect = 'function (...) return 42 end',
+}
+
+format_test {
+    name   = 'Embedding loaded file',
+    single = true,
+    input  = loadfile 'init.lua',
+    expect = 'function (...) ... end',
+}
+
 format_test {
     name = 'Embedding nested function, when on same line is too hard 1',
     single = true,
diff --git a/test/test_pretty.lua b/test/test_pretty.lua
index bce7f0c..f070747 100644
--- a/test/test_pretty.lua
+++ b/test/test_pretty.lua
@@ -341,12 +341,6 @@ SUITE:addTest('UseCase: Can load function from file that is shortly deleted', fu
     assert(true)
 end)
 
-SUITE:addTest('UseCase: Can use pretty on loadstring block', function ()
-    -- TODO: Move to more appropriate test module.
-    format(loadstring 'hi = 0')
-    assert(true)
-end)
-
 --------------------------------------------------------------------------------
 
 return SUITE