Fix multi-line token adjustment

Author: ionite34

StabilityMatrix.Avalonia/Behaviors/TextEditorWeightAdjustmentBehavior.cs

@@ -125,15 +125,25 @@ private void HandleWeightAdjustment(double delta)
 
         try
         {
+            // 1. Tokenize the entire text
+            var text = textEditor.Document.Text;
+            var tokenizeResult = TokenizerProvider.TokenizeLine(text);
+
+            // 2. Get the token segment
             if (
-                (editorSelectionSegment != null ? GetSelectedTokenSpan() : GetCaretTokenSpan())
+                (
+                    editorSelectionSegment != null
+                        ? GetSelectedTokenSpan(tokenizeResult)
+                        : GetCaretTokenSpan(tokenizeResult)
+                )
                 is not { } tokenSegment
             )
+            {
+                Logger.Warn("No token segment found");
                 return;
+            }
 
-            // 2. Tokenize the entire text
-            var text = textEditor.Document.Text;
-            var tokenizeResult = TokenizerProvider.TokenizeLine(text);
+            Logger.Debug("Token segment: {Segment}", tokenSegment);
 
             // 3. Build the AST
             var astBuilder = new PromptSyntaxBuilder(tokenizeResult, text);
@@ -145,7 +155,8 @@ is not { } tokenSegment
             // If empty use intersection instead
             if (selectedNodes.Count == 0)
             {
-                selectedNodes = ast.RootNode.Content.Where(node => node.Span.IntersectsWith(tokenSegment))
+                selectedNodes = ast
+                    .RootNode.Content.Where(node => node.Span.IntersectsWith(tokenSegment))
                     .ToList();
             }
 
@@ -174,7 +185,11 @@ is not { } tokenSegment
 
             // Go up and find the first parenthesized node, if any
             // (Considering only the first and last of the smallest nodes)
-            var parenthesisTargets = smallestNodes.Take(1).Concat(smallestNodes.TakeLast(1)).ToList();
+            var parenthesisTargets = smallestNodes.Take(1).ToList();
+            if (smallestNodes.Count > 1)
+            {
+                parenthesisTargets.Add(smallestNodes.Last());
+            }
 
             Logger.Trace("Parenthesis targets: {Nodes}", parenthesisTargets);
 
@@ -183,6 +198,8 @@ is not { } tokenSegment
                 .OfType<ParenthesizedNode>()
                 .FirstOrDefault();
 
+            // Logger.Trace("Parenthesized node: {Node} of {Nodes}", parenthesizedNode, parenthesisTargets);
+
             var currentWeight = 1.0;
             int replacementOffset; // Offset to start replacing text
             int replacementLength; // Length of text to replace (0 if inserting)
@@ -224,6 +241,12 @@ is not { } tokenSegment
             }
 
             // 8. Replace the text.
+            Logger.Debug(
+                "Replacing source text {SrcText} at {SrcRange}, with new text {NewText}",
+                text[replacementOffset..(replacementOffset + replacementLength)],
+                new TextSpan(replacementOffset, replacementLength),
+                newText
+            );
             textEditor.Document.Replace(replacementOffset, replacementLength, newText);
 
             // Plus 1 to caret if we added parenthesis
@@ -256,7 +279,7 @@ is not { } tokenSegment
         }
     }
 
-    private TextSpan? GetSelectedTokenSpan()
+    private TextSpan? GetSelectedTokenSpan(ITokenizeLineResult result)
     {
         if (textEditor is null)
             return null;
@@ -267,25 +290,14 @@ is not { } tokenSegment
         var selectionStart = textEditor.SelectionStart;
         var selectionEnd = selectionStart + textEditor.SelectionLength;
 
-        var startLine = textEditor.Document.GetLineByOffset(selectionStart);
-        var endLine = textEditor.Document.GetLineByOffset(selectionEnd);
-
-        // For simplicity, we'll assume the selection is on a single line.
-        // Multi-line weight adjustment would be significantly more complex.
-        if (startLine.LineNumber != endLine.LineNumber)
-            return null;
-
-        var lineText = textEditor.Document.GetText(startLine.Offset, startLine.Length);
-        var result = TokenizerProvider!.TokenizeLine(lineText);
-
         IToken? startToken = null;
         IToken? endToken = null;
 
         // Find the tokens that intersect the selection.
         foreach (var token in result.Tokens)
         {
-            var tokenStart = token.StartIndex + startLine.Offset;
-            var tokenEnd = token.EndIndex + startLine.Offset;
+            var tokenStart = token.StartIndex;
+            var tokenEnd = token.EndIndex;
 
             if (tokenEnd > selectionStart && startToken is null)
             {
@@ -306,23 +318,15 @@ is not { } tokenSegment
             return null;
 
         // Ensure end index is within length of text
-        var endIndex = Math.Min(endToken.EndIndex + startLine.Offset, textEditor.Document.TextLength);
+        var endIndex = Math.Min(endToken.EndIndex, textEditor.Document.TextLength);
 
         return TextSpan.FromBounds(startToken.StartIndex, endIndex);
     }
 
-    private TextSpan? GetCaretTokenSpan()
+    private TextSpan? GetCaretTokenSpan(ITokenizeLineResult result)
     {
-        var caret = textEditor!.CaretOffset;
-
-        // Get the line the caret is on
-        var line = textEditor.Document.GetLineByOffset(caret);
-        var lineText = textEditor.Document.GetText(line.Offset, line.Length);
-
-        var caretAbsoluteOffset = caret - line.Offset;
-
-        // Tokenize
-        var result = TokenizerProvider!.TokenizeLine(lineText);
+        var caretAbsoluteOffset = textEditor!.CaretOffset;
+        var textEndOffset = textEditor.Document.TextLength;
 
         IToken? currentToken = null;
         var currentTokenIndex = -1;
@@ -331,11 +335,11 @@ is not { } tokenSegment
         {
             var token = result.Tokens[i];
             // If we see a line comment token anywhere, return null
-            var isComment = token.Scopes.Any(s => s.Contains("comment.line"));
+            /*var isComment = token.Scopes.Any(s => s.Contains("comment.line"));
             if (isComment)
             {
                 return null;
-            }
+            }*/
 
             // Find match
             if (caretAbsoluteOffset >= token.StartIndex && caretAbsoluteOffset < token.EndIndex)
@@ -354,18 +358,33 @@ is not { } tokenSegment
             }
         }
 
-        // Check if the token is a separator, if so check the next token instead
-        if (
-            currentToken?.Scopes is { } scopes
-            && scopes.Contains("meta.structure.array.prompt")
-            && scopes.Contains("punctuation.separator.variable.prompt")
-        )
+        // Check if the token is a separator, if so check the previous or next token instead
+        if (currentToken?.Scopes is { } scopes && scopes.Contains("punctuation.separator.variable.prompt"))
         {
             // Check if we have a prev token
-            var nextToken = result.Tokens.ElementAtOrDefault(currentTokenIndex - 1);
-            if (nextToken is not null)
+            if (
+                result.Tokens.ElementAtOrDefault(currentTokenIndex - 1) is { } prevToken
+                && !prevToken.Scopes.Contains("punctuation.separator.variable.prompt")
+            )
+            {
+                Logger.Trace(
+                    "Matched on seperator, using previous token: {Current} -> {Prev}",
+                    currentToken,
+                    prevToken
+                );
+                currentToken = prevToken;
+            }
+            // Check if we have a next token
+            else if (
+                result.Tokens.ElementAtOrDefault(currentTokenIndex + 1) is { } nextToken
+                && !nextToken.Scopes.Contains("punctuation.separator.variable.prompt")
+            )
             {
-                // Use the next token instead
+                Logger.Trace(
+                    "Matched on seperator, using next token: {Current} -> {Next}",
+                    currentToken,
+                    nextToken
+                );
                 currentToken = nextToken;
             }
         }
@@ -382,11 +401,11 @@ is not { } tokenSegment
             return null;
         }
 
-        var startOffset = currentToken.StartIndex + line.Offset;
-        var endOffset = currentToken.EndIndex + line.Offset;
-
         // Cap the offsets by the line offsets
-        return TextSpan.FromBounds(Math.Max(startOffset, line.Offset), Math.Min(endOffset, line.EndOffset));
+        var startOffset = Math.Max(currentToken.StartIndex, 0);
+        var endOffset = Math.Min(currentToken.EndIndex, textEndOffset);
+
+        return TextSpan.FromBounds(startOffset, endOffset);
     }
 
     [Localizable(false)]

StabilityMatrix.Core/Models/PromptSyntax/PromptSyntaxBuilder.cs

@@ -28,7 +28,7 @@ public PromptSyntaxTree BuildAST()
         var rootNode = new DocumentNode
         {
             Span = new TextSpan(startIndex, endIndex - startIndex),
-            Content = nodes
+            Content = nodes,
         };
 
         return new PromptSyntaxTree(sourceText, rootNode, tokenizeResult.Tokens.ToList());
@@ -163,7 +163,7 @@ private NumberNode ParseNumber()
         {
             Raw = number,
             Span = new TextSpan(token.StartIndex, token.Length),
-            Value = decimal.Parse(number, CultureInfo.InvariantCulture)
+            Value = decimal.Parse(number, CultureInfo.InvariantCulture),
         };
     }
 
@@ -181,8 +181,14 @@ openParenToken is null
 
         while (MoreTokens())
         {
+            // Check if no more tokens to consume.
             if (PeekToken() is not { } nextToken)
+            {
+                // Ensure we have length set
+                if (node.Span.Length == 0)
+                    throw new InvalidOperationException("Unexpected end of input.");
                 break;
+            }
 
             if (nextToken.Scopes.Contains("punctuation.separator.weight.prompt"))
             {
@@ -200,10 +206,17 @@ openParenToken is null
                 node.Weight = ParseNumber();
             }
             // We're supposed to check `punctuation.definition.array.end.prompt` here, textmate is not parsing it
-            // separately with current tmLanguage grammar, so use `meta.structure.weight.prompt` for now
+            // separately always with current tmLanguage grammar, so ALSO use `meta.structure.weight.prompt` for now
             // We check this AFTER `punctuation.separator.weight.prompt` to avoid consuming the ':'
-            else if (nextToken.Scopes.Contains("meta.structure.weight.prompt"))
+            else if (
+                nextToken.Scopes.Contains("punctuation.definition.array.end.prompt")
+                || nextToken.Scopes.Contains("meta.structure.weight.prompt")
+            )
             {
+                // Verify contents
+                if (GetTextSubstring(nextToken) != ")")
+                    throw new InvalidOperationException("Expected closing parenthesis.");
+
                 ConsumeToken(); // Consume the ')'
                 node.EndIndex = nextToken.EndIndex; // Set end index
                 break;
@@ -287,7 +300,7 @@ endNetworkToken is null
             NetworkType = type,
             ModelName = name,
             ModelWeight = modelWeight,
-            ClipWeight = clipWeight
+            ClipWeight = clipWeight,
         };
     }
 
@@ -299,7 +312,7 @@ private ArrayNode ParseArray()
 
         var node = new ArrayNode
         {
-            Span = new TextSpan(openBracket.StartIndex, 0) // Set start index
+            Span = new TextSpan(openBracket.StartIndex, 0), // Set start index
         };
 
         while (MoreTokens())
@@ -334,7 +347,7 @@ openBraceToken is null
 
         var node = new WildcardNode
         {
-            Span = new TextSpan(openBraceToken.StartIndex, 0) // Set start index
+            Span = new TextSpan(openBraceToken.StartIndex, 0), // Set start index
         };
 
         while (MoreTokens())
@@ -378,7 +391,7 @@ openBraceToken is null
                         StartIndex = result.StartIndex,
                         EndIndex = sourceText.Length,
                         Length = sourceText.Length - result.StartIndex,
-                        Scopes = result.Scopes
+                        Scopes = result.Scopes,
                     };
                 }
             }
@@ -405,7 +418,7 @@ private IToken ConsumeToken()
                     StartIndex = result.StartIndex,
                     EndIndex = sourceText.Length,
                     Length = sourceText.Length - result.StartIndex,
-                    Scopes = result.Scopes
+                    Scopes = result.Scopes,
                 };
             }
         }