From 400d1d927d76791b46ae30d431d908c60e411a26 Mon Sep 17 00:00:00 2001
From: Jeremy Fleischman <jeremyfleischman@gmail.com>
Date: Wed, 18 Sep 2024 14:41:21 -0700
Subject: [PATCH] lib: fix escaping bugs in wrapVimscriptForLua and
 wrapLuaForVimscript

These functions had very similar bugs: they didn't check if their chosen
"close token" was already present in the string they're escaping.

I went ahead and did the work implied by the TODOs: search for a "close
token" that is *not* in the original string. Pretty simple concept, but
it turned into an annoying amount of code. I couldn't find anything in
upstream nixpkgs lib, or some clever insight about lua/vimscript that
makes this work unecessary, but I'll be thrilled (and a little bummed
about a wasted afternoon) to learn about something.
---
 lib/utils.nix       | 99 ++++++++++++++++++++++++++++++++++++++-------
 tests/lib-tests.nix | 35 ++++++++++++++++
 2 files changed, 120 insertions(+), 14 deletions(-)

diff --git a/lib/utils.nix b/lib/utils.nix
index 72558c88..f36b1b70 100644
--- a/lib/utils.nix
+++ b/lib/utils.nix
@@ -125,27 +125,98 @@ rec {
     end
   '';
 
+  /**
+    Convert the given String to a Lua [long literal].
+    For example, you could use this to safely pass a Vimscript string to the
+    `vim.cmd` function.
+
+    [long literal]: https://www.lua.org/manual/5.4/manual.html#3.1
+
+    # Examples
+
+    ```nix
+    nix-repl> toLuaLongLiteral "simple"
+    "[[simple]]"
+    ```
+
+    ```nix
+    nix-repl> toLuaLongLiteral "]]"
+    "[=[]]]=]"
+    ```
+
+    # Type
+
+    ```
+    toLuaLongLiteral :: String -> String
+    ```
+  */
+  toLuaLongLiteral =
+    string:
+    let
+      findTokens =
+        depth:
+        let
+          infix = lib.strings.replicate depth "=";
+          tokens.open = "[${infix}[";
+          tokens.close = "]${infix}]";
+        in
+        if lib.hasInfix tokens.close string then findTokens (depth + 1) else tokens;
+
+      tokens = findTokens 0;
+    in
+    tokens.open + string + tokens.close;
+
+  /**
+    Convert the given String into a Vimscript [:let-heredoc].
+    For example, you could use this to invoke [:lua].
+
+    [:let-heredoc]: https://neovim.io/doc/user/eval.html#%3Alet-heredoc
+    [:lua]: https://neovim.io/doc/user/lua.html#%3Alua-heredoc
+
+    # Examples
+
+    ```nix
+    toVimscriptHeredoc "simple"
+    => "<< EOF\nsimple\nEOF"
+    ```
+
+    ```nix
+    toVimscriptHeredoc "EOF"
+    => "<< EOFF\nEOF\nEOFF"
+    ```
+
+    # Type
+
+    ```
+    toVimscriptHeredoc :: String -> String
+    ```
+  */
+  toVimscriptHeredoc =
+    string:
+    let
+      findToken =
+        depth:
+        let
+          token = "EOF" + lib.strings.replicate depth "F";
+        in
+        if lib.hasInfix token string then findToken (depth + 1) else token;
+
+      token = findToken 0;
+    in
+    ''
+      << ${token}
+      ${string}
+      ${token}'';
+
   # Wrap Vimscript for using in lua,
   # but only if the string contains something other than whitespaces
-  # TODO: account for a possible ']]' in the string
   wrapVimscriptForLua =
-    string:
-    lib.optionalString (hasContent string) ''
-      vim.cmd([[
-      ${string}
-      ]])
-    '';
+    string: lib.optionalString (hasContent string) "vim.cmd(${toLuaLongLiteral string})";
 
   # Wrap lua script for using in Vimscript,
   # but only if the string contains something other than whitespaces
-  # TODO: account for a possible 'EOF' if the string
   wrapLuaForVimscript =
-    string:
-    lib.optionalString (hasContent string) ''
-      lua << EOF
-      ${string}
-      EOF
-    '';
+    string: lib.optionalString (hasContent string) "lua ${toVimscriptHeredoc string}";
 
   # Split a list into a several sub-list, each with a max-size of `size`
   groupListBySize =
diff --git a/tests/lib-tests.nix b/tests/lib-tests.nix
index abaecb23..870669c8 100644
--- a/tests/lib-tests.nix
+++ b/tests/lib-tests.nix
@@ -374,6 +374,41 @@ let
         ];
       };
     };
+
+    testEscapeStringForLua = {
+      expr = lib.mapAttrs (_: helpers.utils.toLuaLongLiteral) {
+        simple = "simple";
+        depth-one = " ]] ";
+        depth-two = " ]] ]=] ";
+      };
+      expected = {
+        simple = "[[simple]]";
+        depth-one = "[=[ ]] ]=]";
+        depth-two = "[==[ ]] ]=] ]==]";
+      };
+    };
+
+    testEscapeStringForVimscript = {
+      expr = lib.mapAttrs (_: helpers.utils.toVimscriptHeredoc) {
+        simple = "simple";
+        depth-one = "EOF";
+        depth-two = "EOF EOFF";
+      };
+      expected = {
+        simple = ''
+          << EOF
+          simple
+          EOF'';
+        depth-one = ''
+          << EOFF
+          EOF
+          EOFF'';
+        depth-two = ''
+          << EOFFF
+          EOF EOFF
+          EOFFF'';
+      };
+    };
   };
 in
 if results == [ ] then