Populate ScriptTokenStream and propagate precise token bounds during AST construction

scripts/generate_parser.py

@@ -281,7 +281,7 @@ def main():
         'TryParseOption'
     ]
     void_methods = [
-        'initialize', 'beginComplexToken', 'checkEOF', 'UpdateTokenInfo',
+        'initialize', 'beginComplexToken', 'checkEOF',
         'UpdatePositionInfo', 'AddParseError',
         'ThrowParseErrorException', 'ThrowIncorrectSyntaxErrorException',
         'ThrowPartialAstIfPhaseOne', 'ThrowConstraintIfPhaseOne',
@@ -358,7 +358,10 @@ def main():
         '    public Identifier GetEmptyIdentifier(Object... args) { return new Identifier(); }',
         '    public Literal GetIPv4FragmentFromDotNumberNumeric(Object... args) { return new IntegerLiteral(); }',
         '    public Literal GetIPv4FragmentFromNumberDotNumeric(Object... args) { return new IntegerLiteral(); }',
-        '    public void AddAndUpdateTokenInfo(Object... args) { if (args.length >= 3 && args[1] instanceof java.util.List list && args[2] != null) { if (args[2] instanceof java.util.Collection<?> c) { list.addAll(c); } else { list.add(args[2]); } if (args[0] instanceof SchemaObjectName son) { int count = son.Identifiers.size(); if (count > 0) son.BaseIdentifier = son.Identifiers.get(count - 1); if (count > 1) son.SchemaIdentifier = son.Identifiers.get(count - 2); if (count > 2) son.DatabaseIdentifier = son.Identifiers.get(count - 3); if (count > 3) son.ServerIdentifier = son.Identifiers.get(count - 4); } } }',
+        '    public void AddAndUpdateTokenInfo(Object... args) { if (args.length >= 3 && args[1] instanceof java.util.List list && args[2] != null) { if (args[2] instanceof java.util.Collection<?> c) { list.addAll(c); for (Object item : c) { UpdateTokenInfo(args[0], item); } } else { list.add(args[2]); UpdateTokenInfo(args[0], args[2]); } if (args[0] instanceof SchemaObjectName son) { int count = son.Identifiers.size(); if (count > 0) son.BaseIdentifier = son.Identifiers.get(count - 1); if (count > 1) son.SchemaIdentifier = son.Identifiers.get(count - 2); if (count > 2) son.DatabaseIdentifier = son.Identifiers.get(count - 3); if (count > 3) son.ServerIdentifier = son.Identifiers.get(count - 4); } } }',
+        '    public int ResolveTokenIndex(Object tokenLike) { if (_tokenSource == null || tokenLike == null) return -1; if (tokenLike instanceof Token tok) { return _tokenSource._getTokenIndex(tok); } return -1; }',
+        '    public void AttachScriptTokenStream(Object fragment) { if (_tokenSource == null || !(fragment instanceof TSqlFragment f) || f.ScriptTokenStream == null) return; if (f.ScriptTokenStream.isEmpty()) { f.ScriptTokenStream.addAll(_tokenSource._getScriptTokenStream()); } }',
+        '    public void UpdateTokenInfo(Object... args) { if (args == null || args.length == 0 || !(args[0] instanceof TSqlFragment fragment)) return; AttachScriptTokenStream(fragment); if (args.length >= 3 && args[1] instanceof Number first && args[2] instanceof Number last) { fragment.UpdateTokenInfo(first.intValue(), last.intValue()); return; } if (args.length >= 2 && args[1] instanceof TSqlFragment child) { fragment.UpdateTokenInfo(child); return; } if (args.length >= 2) { int idx = ResolveTokenIndex(args[1]); if (idx >= 0) { if (fragment.firstTokenIndex == 0 || idx < fragment.firstTokenIndex) { fragment.firstTokenIndex = idx; } if (idx > fragment.lastTokenIndex) { fragment.lastTokenIndex = idx; } } } }',
         '    public void AddIdentifierToListWithCheck(Object... args) { if (args.length >= 2 && args[0] instanceof java.util.List list && args[1] != null) { if (args[1] instanceof java.util.Collection<?> c) { list.addAll(c); } else { list.add(args[1]); } } }',
         '    public SqlDataTypeOption ParseDataType180(Object s) { if (s == null) return SqlDataTypeOption.None; String str = s.toString(); for (SqlDataTypeOption opt : SqlDataTypeOption.values()) { if (opt.name().equalsIgnoreCase(str)) return opt; } return SqlDataTypeOption.None; }',
         '    public SqlDataTypeOption ParseDataType160(Object s) { return ParseDataType180(s); }',

tsql-scriptdom-astgen/src/main/java/io/github/fupacat/AstCodeGenerator.java

@@ -57,6 +57,14 @@ private int _generateBaseFragment(Path outputDirectory) throws Exception {
                 .addMethod(com.squareup.javapoet.MethodSpec.methodBuilder("UpdateTokenInfo")
                         .addModifiers(javax.lang.model.element.Modifier.PUBLIC)
                         .addParameter(Object.class, "other")
+                        .beginControlFlow("if (other instanceof TSqlFragment fragment)")
+                        .beginControlFlow("if (this.firstTokenIndex == 0 || (fragment.firstTokenIndex > 0 && fragment.firstTokenIndex < this.firstTokenIndex))")
+                        .addStatement("this.firstTokenIndex = fragment.firstTokenIndex")
+                        .endControlFlow()
+                        .beginControlFlow("if (fragment.lastTokenIndex > this.lastTokenIndex)")
+                        .addStatement("this.lastTokenIndex = fragment.lastTokenIndex")
+                        .endControlFlow()
+                        .endControlFlow()
                         .build())
                 .addMethod(com.squareup.javapoet.MethodSpec.methodBuilder("SetUnquotedIdentifier")
                         .addModifiers(javax.lang.model.element.Modifier.PUBLIC)

tsql-scriptdom-core/src/main/java/io/github/fupacat/TSqlWhitespaceTokenFilter.java

@@ -6,6 +6,8 @@
 public class TSqlWhitespaceTokenFilter implements TokenStream {
 
     private final java.util.List<TSqlParserToken> streamToFilter;
+    private final java.util.List<TSqlParserToken> scriptTokenStream = new java.util.ArrayList<>();
+    private final java.util.IdentityHashMap<Token, Integer> tokenIndices = new java.util.IdentityHashMap<>();
     private int currentTokenIndex = 0;
     private boolean quotedIdentifier = true;
     public boolean QuotedIdentifier = true;
@@ -56,10 +58,14 @@ public Token nextToken() throws antlr.TokenStreamException {
                     lastToken.setLine(token.getLine());
                     lastToken.setColumn(token.getColumn());
                 }
+                tokenIndices.put(token, scriptTokenStream.size());
+                scriptTokenStream.add(lastToken);
                 return token;
             }
             Token eofToken = new antlr.CommonToken(TSqlTokenTypes.EOF, "");
             lastToken = new TSqlParserToken(TSqlTokenType.EndOfFile, "");
+            tokenIndices.put(eofToken, scriptTokenStream.size());
+            scriptTokenStream.add(lastToken);
             return eofToken;
         }
         TSqlParserToken token = null;
@@ -94,6 +100,8 @@ public Token nextToken() throws antlr.TokenStreamException {
         }
 
         lastToken = token;
+        tokenIndices.put(token, scriptTokenStream.size());
+        scriptTokenStream.add(token);
         return token;
     }
 
@@ -106,4 +114,15 @@ public boolean _isQuotedIdentifier() {
     public TSqlParserToken _getLastToken() {
         return lastToken;
     }
+
+    // TRACEABILITY_MAPPING: REQ-007
+    public java.util.List<TSqlParserToken> _getScriptTokenStream() {
+        return scriptTokenStream;
+    }
+
+    // TRACEABILITY_MAPPING: REQ-007
+    public int _getTokenIndex(Token token) {
+        Integer index = tokenIndices.get(token);
+        return index != null ? index : -1;
+    }
 }

tsql-scriptdom-core/src/test/java/io/github/fupacat/TSqlParserFacadeTest.java

@@ -1,9 +1,11 @@
 package io.github.fupacat;
 
 import io.github.fupacat.ast.SelectStatement;
+import io.github.fupacat.ast.TSqlBatch;
 import io.github.fupacat.ast.TSqlConcreteFragmentVisitor;
 import io.github.fupacat.ast.TSqlFragment;
 import io.github.fupacat.ast.TSqlScript;
+import io.github.fupacat.ast.TSqlStatement;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 
@@ -84,6 +86,23 @@ public void testBuilderAndReaderParseScriptString() {
         assertThat(result.isSuccess()).isTrue();
     }
 
+    @Test
+    public void testParsePopulatesScriptTokenStreamAndTokenBounds() {
+        ParseResult<TSqlScript> result = TSqlParser.parse("SELECT 1 AS a; SELECT 2 AS b;");
+
+        assertThat(result.isSuccess()).isTrue();
+        assertThat(result.fragment()).isNotNull();
+        assertThat(result.fragment().ScriptTokenStream).isNotEmpty();
+        assertThat(result.fragment().lastTokenIndex).isEqualTo(result.fragment().ScriptTokenStream.size() - 1);
+
+        TSqlBatch batch = result.fragment().Batches.get(0);
+        assertThat(batch.lastTokenIndex).isGreaterThan(batch.firstTokenIndex);
+
+        TSqlStatement firstStatement = batch.Statements.get(0);
+        assertThat(firstStatement.lastTokenIndex).isGreaterThan(0);
+        assertThat(firstStatement.lastTokenIndex).isGreaterThanOrEqualTo(firstStatement.firstTokenIndex);
+    }
+
     // TRACEABILITY_MAPPING: REQ-028
     @Test
     public void testReaderOverloadsUseDefaultAndTargetVersion() {