1 files changed, 65 insertions, 9 deletions

diff --git a/lua/lexers/markdown.lua b/lua/lexers/markdown.lua
index b6be9f7..d5a18b8 100644
--- a/lua/lexers/markdown.lua
+++ b/lua/lexers/markdown.lua
@@ -29,22 +29,54 @@ lex:add_rule('list', lex:tag(lexer.LIST,
 local hspace = lexer.space - '\n'
 local blank_line = '\n' * hspace^0 * ('\n' + P(-1))
 
-local code_line = lexer.starts_line((B('    ') + B('\t')) * lexer.to_eol(), true)
-local code_block =
-	lexer.range(lexer.starts_line('```', true), '\n```' * hspace^0 * ('\n' + P(-1))) +
-		lexer.range(lexer.starts_line('~~~', true), '\n~~~' * hspace^0 * ('\n' + P(-1)))
+local code_line = lexer.starts_line((B('    ') + B('\t')) * lpeg.P(function(input, index)
+	-- Backtrack to the start of the current paragraph, which is either after a blank line,
+	-- at the start of a higher level of indentation, or at the start of the buffer.
+	local line, blank_line = lexer.line_from_position(index), false
+	while line > 0 do
+		local s, e = lexer.line_start[line], lexer.line_end[line]
+		blank_line = s == e or lexer.text_range(s, e - s + 1):find('^%s+$')
+		if blank_line then break end
+		local indent_amount = lexer.indent_amount[line]
+		line = line - 1
+		if line > 0 and lexer.indent_amount[line] > indent_amount then break end
+	end
+
+	-- If the start of the paragraph does not being with a '    ' or '\t', then this line
+	-- is a continuation of the current paragraph, not a code block.
+	local text = lexer.text_range(lexer.line_start[line + 1], 4)
+	if not text:find('^\t') and text ~= '    ' then return false end
+
+	-- If the current paragraph is a code block, then so is this line.
+	if line <= 1 then return true end
+
+	-- Backtrack to see if this line is in a list item. If so, it is not a code block.
+	while line > 1 do
+		line = line - 1
+		local s, e = lexer.line_start[line], lexer.line_end[line]
+		local blank = s == e or lexer.text_range(s, e - s + 1):find('^%s+$')
+		if not blank and lexer.indent_amount[line] == 0 then break end
+	end
+	text = lexer.text_range(lexer.line_start[line], 8) -- note: only 2 is needed for unordered lists
+	if text:find('^[*+-][ \t]') then return false end
+	if text:find('^%d+%.[ \t]') then return false end
+
+	return true -- if all else fails, it is probably a code block
+end) * lexer.to_eol(), true)
+
+local code_block = lexer.range(lexer.starts_line('```', true),
+	'\n' * hspace^0 * '```' * hspace^0 * ('\n' + P(-1))) +
+	lexer.range(lexer.starts_line('~~~', true), '\n' * hspace^0 * '~~~' * hspace^0 * ('\n' + P(-1)))
+
 local code_inline = lpeg.Cmt(lpeg.C(P('`')^1), function(input, index, bt)
 	-- `foo`, ``foo``, ``foo`bar``, `foo``bar` are all allowed.
 	local _, e = input:find('[^`]' .. bt .. '%f[^`]', index)
 	return (e or #input) + 1
 end)
+
 lex:add_rule('block_code', lex:tag(lexer.CODE, code_line + code_block + code_inline))
 
-lex:add_rule('blockquote',
-	lex:tag(lexer.STRING, lpeg.Cmt(lexer.starts_line('>', true), function(input, index)
-		local _, e = input:find('\n[ \t]*\r?\n', index) -- the next blank line (possibly with indentation)
-		return (e or #input) + 1
-	end)))
+lex:add_rule('blockquote', lex:tag(lexer.STRING, lexer.starts_line('>', true)))
 
 -- Span elements.
 lex:add_rule('escape', lex:tag(lexer.DEFAULT, P('\\') * 1))
@@ -92,4 +124,28 @@ local start_rule = lexer.starts_line(P(' ')^-3) * #P('<') * html:get_rule('tag')
 local end_rule = #blank_line * ws
 lex:embed(html, start_rule, end_rule)
 
+local FOLD_HEADER, FOLD_BASE = lexer.FOLD_HEADER, lexer.FOLD_BASE
+-- Fold '#' headers.
+function lex:fold(text, start_line, start_level)
+	local levels = {}
+	local line_num = start_line
+	if start_level > FOLD_HEADER then start_level = start_level - FOLD_HEADER end
+	for line in (text .. '\n'):gmatch('(.-)\r?\n') do
+		local header = line:match('^%s*(#*)')
+		-- If the previous line was a header, this line's level has been pre-defined.
+		-- Otherwise, use the previous line's level, or if starting to fold, use the start level.
+		local level = levels[line_num] or levels[line_num - 1] or start_level
+		if level > FOLD_HEADER then level = level - FOLD_HEADER end
+		-- If this line is a header, set its level to be one less than the header level
+		-- (so it can be a fold point) and mark it as a fold point.
+		if #header > 0 then
+			level = FOLD_BASE + #header - 1 + FOLD_HEADER
+			levels[line_num + 1] = FOLD_BASE + #header
+		end
+		levels[line_num] = level
+		line_num = line_num + 1
+	end
+	return levels
+end
+
 return lex