hybridized approach with v1

Author: dzsquared

SqlScriptDom/ScriptDom/SqlServer/ScriptGenerator/SqlScriptGeneratorVisitor.Comments.cs

@@ -23,6 +23,16 @@ internal abstract partial class SqlScriptGeneratorVisitor
         /// </summary>
         private IList<TSqlParserToken> _currentTokenStream;
 
+        /// <summary>
+        /// Tracks which comment tokens have already been emitted to avoid duplicates.
+        /// </summary>
+        private readonly HashSet<TSqlParserToken> _emittedComments = new HashSet<TSqlParserToken>();
+
+        /// <summary>
+        /// Tracks whether leading (file-level) comments have been emitted.
+        /// </summary>
+        private bool _leadingCommentsEmitted = false;
+
         #endregion
 
         #region Comment Preservation Methods
@@ -36,216 +46,226 @@ protected void SetTokenStreamForComments(IList<TSqlParserToken> tokenStream)
         {
             _currentTokenStream = tokenStream;
             _lastProcessedTokenIndex = -1;
+            _emittedComments.Clear();
+            _leadingCommentsEmitted = false;
         }
 
         /// <summary>
-        /// Gets leading comments that appear between the last processed token and the current fragment.
+        /// Emits comments that appear before the first fragment in the script (file-level leading comments).
+        /// Called once when generating the first fragment.
         /// </summary>
-        /// <param name="fragment">The current fragment being visited.</param>
-        /// <returns>List of comment information for leading comments.</returns>
-        protected List<CommentInfo> GetLeadingComments(TSqlFragment fragment)
+        /// <param name="fragment">The first fragment being generated.</param>
+        protected void EmitLeadingComments(TSqlFragment fragment)
         {
-            var comments = new List<CommentInfo>();
-
-            if (_currentTokenStream == null || fragment == null || !_options.PreserveComments)
+            if (!_options.PreserveComments || _currentTokenStream == null || fragment == null)
             {
-                return comments;
+                return;
             }
 
-            int startIndex = _lastProcessedTokenIndex + 1;
-            int endIndex = fragment.FirstTokenIndex;
+            if (fragment.FirstTokenIndex <= 0)
+            {
+                return;
+            }
 
-            // Scan for comments between last processed and current fragment
-            for (int i = startIndex; i < endIndex && i < _currentTokenStream.Count; i++)
+            for (int i = 0; i < fragment.FirstTokenIndex && i < _currentTokenStream.Count; i++)
             {
                 var token = _currentTokenStream[i];
-                if (IsCommentToken(token))
+                if (IsCommentToken(token) && !_emittedComments.Contains(token))
                 {
-                    comments.Add(new CommentInfo(token, CommentPosition.Leading, fragment.FirstTokenIndex));
+                    EmitCommentToken(token, isLeading: true);
+                    _emittedComments.Add(token);
                 }
             }
-
-            return comments;
         }
 
         /// <summary>
-        /// Gets trailing comments that appear on the same line after the fragment.
+        /// Emits comments that appear in the gap between the last emitted token and the current fragment.
+        /// This captures comments embedded within sub-expressions.
         /// </summary>
-        /// <param name="fragment">The current fragment being visited.</param>
-        /// <returns>List of comment information for trailing comments.</returns>
-        protected List<CommentInfo> GetTrailingComments(TSqlFragment fragment)
+        /// <param name="fragment">The fragment about to be generated.</param>
+        protected void EmitGapComments(TSqlFragment fragment)
         {
-            var comments = new List<CommentInfo>();
-
-            if (_currentTokenStream == null || fragment == null || !_options.PreserveComments)
+            if (!_options.PreserveComments || _currentTokenStream == null || fragment == null)
             {
-                return comments;
+                return;
             }
 
-            int lastTokenIndex = fragment.LastTokenIndex;
-            if (lastTokenIndex < 0 || lastTokenIndex >= _currentTokenStream.Count)
+            int startIndex = _lastProcessedTokenIndex + 1;
+            int endIndex = fragment.FirstTokenIndex;
+
+            if (endIndex <= startIndex)
             {
-                return comments;
+                return;
             }
 
-            var lastToken = _currentTokenStream[lastTokenIndex];
-            int lastTokenLine = lastToken.Line;
-
-            // Scan for comments after the last token on the same line
-            for (int i = lastTokenIndex + 1; i < _currentTokenStream.Count; i++)
+            for (int i = startIndex; i < endIndex && i < _currentTokenStream.Count; i++)
             {
                 var token = _currentTokenStream[i];
-                
-                // Stop if we've gone past the same line (unless it's whitespace with no newline)
-                if (token.Line > lastTokenLine)
-                {
-                    break;
-                }
-
-                // Found a comment on the same line
-                if (IsCommentToken(token))
+                if (IsCommentToken(token) && !_emittedComments.Contains(token))
                 {
-                    comments.Add(new CommentInfo(token, CommentPosition.Trailing, lastTokenIndex));
-                    
-                    // For single-line comments, there can't be anything else after on this line
-                    if (token.TokenType == TSqlTokenType.SingleLineComment)
-                    {
-                        break;
-                    }
+                    EmitCommentToken(token, isLeading: true);
+                    _emittedComments.Add(token);
+                    _lastProcessedTokenIndex = i;
                 }
             }
-
-            return comments;
         }
 
         /// <summary>
-        /// Emits a comment using the script writer with current indentation.
+        /// Emits trailing comments that appear immediately after the fragment.
         /// </summary>
-        /// <param name="commentInfo">The comment to emit.</param>
-        protected void EmitComment(CommentInfo commentInfo)
+        /// <param name="fragment">The fragment that was just generated.</param>
+        protected void EmitTrailingComments(TSqlFragment fragment)
         {
-            if (commentInfo == null || commentInfo.Token == null)
+            if (!_options.PreserveComments || _currentTokenStream == null || fragment == null)
             {
                 return;
             }
 
-            var text = commentInfo.Text;
-            if (string.IsNullOrEmpty(text))
+            int lastTokenIndex = fragment.LastTokenIndex;
+            if (lastTokenIndex < 0 || lastTokenIndex >= _currentTokenStream.Count)
             {
                 return;
             }
 
-            if (commentInfo.IsSingleLineComment)
-            {
-                EmitSingleLineComment(text, commentInfo.Position);
-            }
-            else if (commentInfo.IsMultiLineComment)
+            // Scan for comments immediately following the fragment
+            for (int i = lastTokenIndex + 1; i < _currentTokenStream.Count; i++)
             {
-                EmitMultiLineComment(text, commentInfo.Position);
+                var token = _currentTokenStream[i];
+                
+                if (IsCommentToken(token) && !_emittedComments.Contains(token))
+                {
+                    EmitCommentToken(token, isLeading: false);
+                    _emittedComments.Add(token);
+                    _lastProcessedTokenIndex = i;
+                }
+                else if (token.TokenType != TSqlTokenType.WhiteSpace)
+                {
+                    // Stop at next non-whitespace, non-comment token
+                    break;
+                }
             }
         }
 
         /// <summary>
-        /// Emits a single-line comment.
+        /// Updates tracking after generating a fragment.
         /// </summary>
-        private void EmitSingleLineComment(string text, CommentPosition position)
+        /// <param name="fragment">The fragment that was just generated.</param>
+        protected void UpdateLastProcessedIndex(TSqlFragment fragment)
         {
-            if (position == CommentPosition.Trailing)
+            if (fragment != null && fragment.LastTokenIndex > _lastProcessedTokenIndex)
             {
-                // Trailing: add space before comment, keep on same line
-                _writer.AddToken(ScriptGeneratorSupporter.CreateWhitespaceToken(1));
+                _lastProcessedTokenIndex = fragment.LastTokenIndex;
             }
-            else
+        }
+
+        /// <summary>
+        /// Called from GenerateFragmentIfNotNull to handle comments before generating a fragment.
+        /// This is the key integration point that enables comments within sub-expressions.
+        /// </summary>
+        /// <param name="fragment">The fragment about to be generated.</param>
+        protected void BeforeGenerateFragment(TSqlFragment fragment)
+        {
+            if (!_options.PreserveComments || _currentTokenStream == null || fragment == null)
             {
-                // Leading: comment goes on its own line
-                // Indentation is already applied by current context
+                return;
             }
 
-            // Write the comment as-is (preserving the -- prefix)
-            _writer.AddToken(new TSqlParserToken(TSqlTokenType.SingleLineComment, text));
-
-            if (position == CommentPosition.Leading)
+            // Emit file-level leading comments once
+            if (!_leadingCommentsEmitted)
             {
-                // After a leading comment, we need a newline
-                _writer.NewLine();
+                EmitLeadingComments(fragment);
+                _leadingCommentsEmitted = true;
             }
+
+            // Emit any comments in the gap between last processed token and this fragment
+            EmitGapComments(fragment);
         }
 
         /// <summary>
-        /// Emits a multi-line comment, preserving internal structure.
+        /// Called from GenerateFragmentIfNotNull to handle comments after generating a fragment.
         /// </summary>
-        private void EmitMultiLineComment(string text, CommentPosition position)
+        /// <param name="fragment">The fragment that was just generated.</param>
+        protected void AfterGenerateFragment(TSqlFragment fragment)
         {
-            if (position == CommentPosition.Trailing)
+            if (!_options.PreserveComments || _currentTokenStream == null || fragment == null)
             {
-                // Trailing: add space before comment
-                _writer.AddToken(ScriptGeneratorSupporter.CreateWhitespaceToken(1));
+                return;
             }
 
-            // For multi-line comments, we preserve the content as-is
-            // The comment includes /* and */ delimiters
-            _writer.AddToken(new TSqlParserToken(TSqlTokenType.MultilineComment, text));
-
-            if (position == CommentPosition.Leading)
-            {
-                // After a leading multi-line comment, add newline
-                _writer.NewLine();
-            }
+            // Emit trailing comments and update tracking
+            EmitTrailingComments(fragment);
+            UpdateLastProcessedIndex(fragment);
         }
 
         /// <summary>
-        /// Called before visiting a fragment to emit any leading comments.
+        /// Emits a comment token to the output.
         /// </summary>
-        /// <param name="fragment">The fragment about to be visited.</param>
-        protected void BeforeVisitFragment(TSqlFragment fragment)
+        /// <param name="token">The comment token.</param>
+        /// <param name="isLeading">True if this is a leading comment, false for trailing.</param>
+        private void EmitCommentToken(TSqlParserToken token, bool isLeading)
         {
-            if (!_options.PreserveComments || _currentTokenStream == null || fragment == null)
+            if (token == null)
             {
                 return;
             }
 
-            var leadingComments = GetLeadingComments(fragment);
-            foreach (var comment in leadingComments)
+            if (token.TokenType == TSqlTokenType.SingleLineComment)
+            {
+                if (!isLeading)
+                {
+                    // Trailing: add space before comment
+                    _writer.AddToken(ScriptGeneratorSupporter.CreateWhitespaceToken(1));
+                }
+
+                _writer.AddToken(new TSqlParserToken(TSqlTokenType.SingleLineComment, token.Text));
+
+                if (isLeading)
+                {
+                    // After a leading comment, add newline
+                    _writer.NewLine();
+                }
+            }
+            else if (token.TokenType == TSqlTokenType.MultilineComment)
             {
-                EmitComment(comment);
+                if (!isLeading)
+                {
+                    // Trailing: add space before comment
+                    _writer.AddToken(ScriptGeneratorSupporter.CreateWhitespaceToken(1));
+                }
+
+                _writer.AddToken(new TSqlParserToken(TSqlTokenType.MultilineComment, token.Text));
+
+                if (isLeading)
+                {
+                    // After a leading multi-line comment, add newline
+                    _writer.NewLine();
+                }
             }
         }
 
         /// <summary>
-        /// Called after visiting a fragment to emit any trailing comments and update tracking.
+        /// Emits any remaining comments at the end of the token stream.
+        /// Call this after visiting the root fragment to capture comments that appear
+        /// after the last statement (end-of-script comments).
         /// </summary>
-        /// <param name="fragment">The fragment that was just visited.</param>
-        protected void AfterVisitFragment(TSqlFragment fragment)
+        protected void EmitRemainingComments()
         {
-            if (!_options.PreserveComments || _currentTokenStream == null || fragment == null)
+            if (!_options.PreserveComments || _currentTokenStream == null)
             {
                 return;
             }
 
-            var trailingComments = GetTrailingComments(fragment);
-            foreach (var comment in trailingComments)
+            // Scan from the last processed token to the end of the token stream
+            for (int i = _lastProcessedTokenIndex + 1; i < _currentTokenStream.Count; i++)
             {
-                EmitComment(comment);
-            }
-
-            // Update the last processed token index
-            if (fragment.LastTokenIndex >= 0)
-            {
-                // Account for any trailing comments we just emitted
-                int newLastIndex = fragment.LastTokenIndex;
-                foreach (var comment in trailingComments)
+                var token = _currentTokenStream[i];
+                if (IsCommentToken(token) && !_emittedComments.Contains(token))
                 {
-                    // Find the index of this comment token
-                    for (int i = newLastIndex + 1; i < _currentTokenStream.Count; i++)
-                    {
-                        if (_currentTokenStream[i] == comment.Token)
-                        {
-                            newLastIndex = i;
-                            break;
-                        }
-                    }
+                    // End-of-script comments: add newline before, emit comment
+                    _writer.NewLine();
+                    _writer.AddToken(new TSqlParserToken(token.TokenType, token.Text));
+                    _emittedComments.Add(token);
                 }
-                _lastProcessedTokenIndex = newLastIndex;
             }
         }
 

SqlScriptDom/ScriptDom/SqlServer/ScriptGenerator/SqlScriptGeneratorVisitor.TSqlBatch.cs

@@ -11,30 +11,19 @@ partial class SqlScriptGeneratorVisitor
     {
         public override void ExplicitVisit(TSqlBatch node)
         {
-            // Emit leading comments before the batch
-            BeforeVisitFragment(node);
-
+            // Comments are now handled automatically by GenerateFragmentIfNotNull
             foreach (TSqlStatement statement in node.Statements)
             {
-                // Emit leading comments before each statement
-                BeforeVisitFragment(statement);
-
                 GenerateFragmentIfNotNull(statement);
 
                 GenerateSemiColonWhenNecessary(statement);
 
-                // Emit trailing comments after each statement
-                AfterVisitFragment(statement);
-
                 if (statement is TSqlStatementSnippet == false)
                 {
                     NewLine();
                     NewLine();
                 }
             }
-
-            // Emit trailing comments after the batch
-            AfterVisitFragment(node);
         }
     }
 }