1 files changed, 39 insertions, 4 deletions
diff --git a/scripts/format-queries.lua b/scripts/format-queries.lua
index 41ab16ce3..ec2197956 100755
--- a/scripts/format-queries.lua
+++ b/scripts/format-queries.lua
@@ -34,6 +34,7 @@ end)
 
 --- Control the indent here. Change to \t if uses tab instead
 local indent_str = "  "
+local textwidth = 100
 
 -- Query to control the formatter
 local format_queries = [[
@@ -174,8 +175,8 @@ local format_queries = [[
   [
     "_"
     name: (identifier)
-  ]
-  (_) @format.cancel-append
+    (_)
+  ] @format.cancel-append
   .
   ")"
   (#not-has-type? @format.cancel-append comment))
@@ -254,9 +255,26 @@ local format_queries = [[
   .
   (capture))
 
+; Separate this query to avoid capture duplication
+(predicate
+  "(" @format.indent.begin @format.cancel-append)
+(predicate
+  (parameters
+    (comment) @format.prepend-newline
+    .
+    (_) @format.cancel-prepend)
+  (#is-start-of-line? @format.prepend-newline))
 (predicate
   (parameters
-    (_) @format.prepend-space))
+    (_) @format.prepend-space)
+  (#set! conditional-newline))
+(predicate
+  (parameters
+    .
+    (capture)
+    . (_) @format.prepend-space)
+  (#set! lookahead-newline)
+  (#set! conditional-newline))
 ;; Workaround to keep the string's content
 (string) @format.keep
 
@@ -322,7 +340,24 @@ local function iter(bufnr, node, lines, q, level)
         if q["format.prepend-newline"][id] then
           lines[#lines + 1] = string.rep(indent_str, level)
         elseif q["format.prepend-space"][id] then
-          lines[#lines] = lines[#lines] .. " "
+          if not q["format.prepend-space"][id]["conditional-newline"] then
+            lines[#lines] = lines[#lines] .. " "
+          elseif child:byte_length() + 1 + #lines[#lines] > textwidth then
+            lines[#lines + 1] = string.rep(indent_str, level)
+          else
+            -- Do a rough guess of the actual byte length. If it's larger than `columns` then add a newline first
+            -- column - byte_end + byte_start
+            local _, _, byte_start = child:start()
+            local _, _, byte_end = node:end_()
+            if
+              q["format.prepend-space"][id]["lookahead-newline"]
+              and textwidth - (byte_end - byte_start) - #lines[#lines] < 0
+            then
+              lines[#lines + 1] = string.rep(indent_str, level)
+            else
+              lines[#lines] = lines[#lines] .. " "
+            end
+          end
         end
       end
       if q["format.replace"][id] then