1 files changed, 17 insertions, 14 deletions
diff --git a/lib/Lex/Preprocessor.cpp b/lib/Lex/Preprocessor.cpp
index 26bb3a9..5689baa 100644
--- a/lib/Lex/Preprocessor.cpp
+++ b/lib/Lex/Preprocessor.cpp
@@ -429,21 +429,17 @@ SourceLocation Preprocessor::AdvanceToTokenCharacter(SourceLocation TokStart,
   return TokStart.getFileLocWithOffset(PhysOffset);
 }
 
-/// \brief Computes the source location just past the end of the
-/// token at this source location.
-///
-/// This routine can be used to produce a source location that
-/// points just past the end of the token referenced by \p Loc, and
-/// is generally used when a diagnostic needs to point just after a
-/// token where it expected something different that it received. If
-/// the returned source location would not be meaningful (e.g., if
-/// it points into a macro), this routine returns an invalid
-/// source location.
-SourceLocation Preprocessor::getLocForEndOfToken(SourceLocation Loc) {
+SourceLocation Preprocessor::getLocForEndOfToken(SourceLocation Loc, 
+                                                 unsigned Offset) {
   if (Loc.isInvalid() || !Loc.isFileID())
     return SourceLocation();
 
   unsigned Len = Lexer::MeasureTokenLength(Loc, getSourceManager(), Features);
+  if (Len > Offset)
+    Len = Len - Offset;
+  else
+    return Loc;
+  
   return AdvanceToTokenCharacter(Loc, Len);
 }
 
@@ -583,11 +579,18 @@ void Preprocessor::RemoveCommentHandler(CommentHandler *Handler) {
   CommentHandlers.erase(Pos);
 }
 
-void Preprocessor::HandleComment(SourceRange Comment) {
+bool Preprocessor::HandleComment(Token &result, SourceRange Comment) {
+  bool AnyPendingTokens = false;
   for (std::vector<CommentHandler *>::iterator H = CommentHandlers.begin(),
        HEnd = CommentHandlers.end();
-       H != HEnd; ++H)
-    (*H)->HandleComment(*this, Comment);
+       H != HEnd; ++H) {
+    if ((*H)->HandleComment(*this, Comment))
+      AnyPendingTokens = true;
+  }
+  if (!AnyPendingTokens || getCommentRetentionState())
+    return false;
+  Lex(result);
+  return true;
 }
 
 CommentHandler::~CommentHandler() { }