format: Fix unicode escapes in string literals

compiler/src/Data/Utf8.hs

@@ -24,6 +24,7 @@ module Data.Utf8
     --
     toChars,
     toText,
+    toShortByteString,
     toBuilder,
     toEscapedBuilder,
     --
@@ -45,6 +46,7 @@ import Data.Binary.Put (putBuilder)
 import Data.Bits (shiftR, (.&.))
 import Data.ByteString.Builder.Internal qualified as B
 import Data.ByteString.Internal qualified as B
+import Data.ByteString.Short qualified as BSS
 import Data.Char qualified as Char
 import Data.List qualified as List
 import Data.Text qualified as Text
@@ -339,6 +341,12 @@ toText =
   -- This could most certainly be optimized for better performance
   Text.pack . toChars
 
+-- TO BYTESTRING
+
+toShortByteString :: Utf8 t -> BSS.ShortByteString
+toShortByteString (Utf8 bytes) =
+  BSS.SBS bytes
+
 -- TO BUILDER
 
 toBuilder :: Utf8 t -> B.Builder

compiler/src/Gren/Format.hs

@@ -13,6 +13,7 @@ import AST.Utils.Binop qualified as Binop
 import Control.Monad (join)
 import Data.Bifunctor (second)
 import Data.ByteString.Builder qualified as B
+import Data.ByteString.Short qualified as BSS
 import Data.List qualified as List
 import Data.List.NonEmpty (NonEmpty ((:|)), nonEmpty)
 import Data.List.NonEmpty qualified as NonEmpty
@@ -22,8 +23,6 @@ import Data.Maybe (catMaybes, maybeToList)
 import Data.Maybe qualified as Maybe
 import Data.Name (Name)
 import Data.Semigroup (sconcat)
-import Data.Text qualified as Text
-import Data.Text.Encoding (encodeUtf8Builder)
 import Data.Utf8 qualified as Utf8
 import Gren.Int qualified as GI
 import Gren.String qualified as GS
@@ -1065,10 +1064,41 @@ formatString style str =
         NonEmpty.fromList $
           mconcat
             [ [Block.line (Block.string7 "\"\"\"")],
-              fmap (Block.line . Block.lineFromBuilder . encodeUtf8Builder) $ Text.splitOn "\\n" $ (Utf8.toText str),
+              fmap (Block.line . Block.lineFromBuilder) $ escapeString True str,
               [Block.line (Block.string7 "\"\"\"")]
             ]
   where
     stringBox :: Block.Line -> Block
     stringBox quotes =
-      Block.line $ quotes <> utf8 str <> quotes
+      Block.line $ quotes <> Block.lineFromBuilder (mconcat $ escapeString False str) <> quotes
+
+    escapeString :: Bool -> Utf8.Utf8 t -> [B.Builder]
+    escapeString isMultiline s =
+      escapeStringHelper isMultiline mempty (Utf8.toShortByteString s)
+
+    escapeStringHelper :: Bool -> B.Builder -> BSS.ShortByteString -> [B.Builder]
+    escapeStringHelper isMultiline acc bytes =
+      case BSS.elemIndex 0x5C {- \ -} bytes of
+        Nothing -> [acc <> B.shortByteString bytes]
+        Just nextBackslashIndex ->
+          case BSS.indexMaybe bytes (nextBackslashIndex + 1) of
+            Just 0x75 {- u -} ->
+              let (pre, hexPlusRest) = BSS.splitAt (nextBackslashIndex + 2) bytes
+                  (hex, rest) = BSS.splitAt 4 hexPlusRest
+               in escapeStringHelper
+                    isMultiline
+                    ( acc
+                        <> B.shortByteString pre
+                        <> B.char7 '{'
+                        <> B.shortByteString hex
+                        <> B.char7 '}'
+                    )
+                    rest
+            Just 0x6E {- n -}
+              | isMultiline ->
+                  let (pre, rest) = BSS.splitAt nextBackslashIndex bytes
+                   in (acc <> B.shortByteString pre)
+                        : escapeStringHelper isMultiline mempty (BSS.drop 2 rest)
+            _ ->
+              let (pre, rest) = BSS.splitAt (nextBackslashIndex + 1) bytes
+               in escapeStringHelper isMultiline (acc <> B.shortByteString pre) rest

tests/Integration/FormatSpec.hs

@@ -350,7 +350,13 @@ spec = do
     describe "string literals" $ do
       it "formats strings" $
         assertFormattedExpression
-          ["a"]
+          ["\"a\""]
+      it "formats escapes" $
+        assertFormattedExpression
+          ["\"\\n\\r\\t\\\"\\'\\\\\""]
+      it "formats unicode escapes" $
+        assertFormattedExpression
+          ["\"A\\u{00A0}B\\u{D83E}\\u{DDDB}C\""]
       it "formats multiline strings with trimmed whitespace" $
         assertFormattedModuleBody
           [ "str =",
@@ -371,6 +377,12 @@ spec = do
                                        "      2",
                                        "    \"\"\""
                                      ]
+      it "formats unicode escapes in multiline strings" $
+        assertFormattedExpression
+          [ "\"\"\"",
+            "A\\u{00A0}B\\u{D83E}\\u{DDDB}C",
+            "\"\"\""
+          ]
 
     describe "int literals" $ do
       it "formats decimal integers" $