4 files changed, 77 insertions, 70 deletions
diff --git a/src/qmlls/qqmlcodemodel_p.h b/src/qmlls/qqmlcodemodel_p.h
index dba6060dff..d2851a885f 100644
--- a/src/qmlls/qqmlcodemodel_p.h
+++ b/src/qmlls/qqmlcodemodel_p.h
@@ -20,6 +20,7 @@
 #include "qtextdocument_p.h"
 #include "qprocessscheduler_p.h"
 #include "qqmllshelputils_p.h"
+#include "qqmlsemantictokens_p.h"
 
 #include <QObject>
 #include <QHash>
@@ -70,7 +71,7 @@ public:
 struct RegisteredSemanticTokens
 {
     QByteArray resultId = "0";
-    QList<int> lastTokens;
+    QmlHighlighting::HighlightsContainer highlights;
 };
 
 struct ModuleSetting
diff --git a/src/qmlls/qqmlhighlightsupport.cpp b/src/qmlls/qqmlhighlightsupport.cpp
index 70f484d245..50617cceb9 100644
--- a/src/qmlls/qqmlhighlightsupport.cpp
+++ b/src/qmlls/qqmlhighlightsupport.cpp
@@ -3,12 +3,14 @@
 // Qt-Security score:significant reason:default
 
 #include <qqmlhighlightsupport_p.h>
+#include <qqmldiffer_p.h>
 
 QT_BEGIN_NAMESPACE
 
 using namespace Qt::StringLiterals;
 using namespace QLspSpecification;
 using namespace QQmlJS::Dom;
+using namespace QmlHighlighting;
 
 /*!
 \internal
@@ -37,7 +39,36 @@ QList<QByteArray> defaultTokenModifiersList()
 
 QList<QByteArray> extendedTokenTypesList()
 {
-    return enumToByteArray<QmlHighlighting::SemanticTokenProtocolTypes>();
+    return enumToByteArray<SemanticTokenProtocolTypes>();
+}
+
+static QList<int> generateHighlights(QmlLsp::RegisteredSemanticTokens &cached,
+                                     const QmlLsp::OpenDocument &doc,
+                                     const std::optional<HighlightsRange> &range,
+                                     HighlightingMode mode)
+{
+    DomItem file = doc.snapshot.doc.fileObject(GoTo::MostLikely);
+    const auto fileObject = file.ownerAs<QmlFile>();
+    QmlHighlighting::Utils::updateResultID(cached.resultId);
+    if (!fileObject || !(fileObject && fileObject->isValid())) {
+        if (const auto lastValidItem = doc.snapshot.validDoc.ownerAs<QmlFile>()) {
+            const auto shiftedHighlights = QmlHighlighting::Utils::shiftHighlights(
+                cached.highlights, lastValidItem->code(), doc.textDocument->toPlainText());
+                return QmlHighlighting::Utils::encodeSemanticTokens(shiftedHighlights, mode);
+        } else {
+            // TODO: Implement regexp based fallback highlighting
+            return {};
+        }
+    } else {
+        HighlightsContainer highlights = QmlHighlighting::Utils::visitTokens(file, range);
+        if (highlights.isEmpty())
+            return {};
+        // Record the highlights for future diffs, only record full highlights
+        if (!range.has_value() )
+            cached.highlights = highlights;
+
+        return QmlHighlighting::Utils::encodeSemanticTokens(highlights, mode);
+    }
 }
 
 /*!
@@ -47,7 +78,7 @@ https://microsoft.github.io/language-server-protocol/specifications/specificatio
 Sends a QLspSpecification::SemanticTokens data as response that is generated for the entire file.
 */
 SemanticTokenFullHandler::SemanticTokenFullHandler(QmlLsp::QQmlCodeModelManager *codeModelManager)
-    : QQmlBaseModule(codeModelManager), m_mode(QmlHighlighting::HighlightingMode::Default)
+    : QQmlBaseModule(codeModelManager), m_mode(HighlightingMode::Default)
 {
 }
 
@@ -63,23 +94,14 @@ void SemanticTokenFullHandler::process(
     ResponseScopeGuard guard(result, request->m_response);
     const QByteArray uri = QQmlLSUtils::lspUriToQmlUrl(request->m_parameters.textDocument.uri);
     const auto doc = m_codeModelManager->openDocumentByUrl(uri);
-    DomItem file = doc.snapshot.doc.fileObject(GoTo::MostLikely);
-    const auto fileObject = file.ownerAs<QmlFile>();
-    if (!fileObject || !(fileObject && fileObject->isValid())) {
-        guard.setError({
-                int(QLspSpecification::ErrorCodes::RequestCancelled),
-                "Cannot proceed: current QML document is invalid!"_L1,
-        });
+    auto &cached = m_codeModelManager->registeredTokens(uri);
+    const auto encoded = generateHighlights(cached, doc, std::nullopt, m_mode);
+
+    if (encoded.isEmpty()) {
+        result = nullptr;
         return;
-    }
-    auto &&encoded = QmlHighlighting::Utils::collectTokens(file, std::nullopt, m_mode);
-    auto &registeredTokens = m_codeModelManager->registeredTokens(uri);
-    if (!encoded.isEmpty()) {
-        QmlHighlighting::Utils::updateResultID(registeredTokens.resultId);
-        result = SemanticTokens{ registeredTokens.resultId, encoded };
-        registeredTokens.lastTokens = std::move(encoded);
     } else {
-        result = nullptr;
+        result = SemanticTokens{cached.resultId, std::move(encoded)};
     }
 }
 
@@ -96,7 +118,7 @@ Sends either SemanticTokens or SemanticTokensDelta data as response.
 This is generally requested when the text document is edited after receiving full highlighting data.
 */
 SemanticTokenDeltaHandler::SemanticTokenDeltaHandler(QmlLsp::QQmlCodeModelManager *codeModelManager)
-    : QQmlBaseModule(codeModelManager), m_mode(QmlHighlighting::HighlightingMode::Default)
+    : QQmlBaseModule(codeModelManager), m_mode(HighlightingMode::Default)
 {
 }
 
@@ -112,33 +134,20 @@ void SemanticTokenDeltaHandler::process(
     ResponseScopeGuard guard(result, request->m_response);
     const QByteArray uri = QQmlLSUtils::lspUriToQmlUrl(request->m_parameters.textDocument.uri);
     const auto doc = m_codeModelManager->openDocumentByUrl(uri);
-    DomItem file = doc.snapshot.doc.fileObject(GoTo::MostLikely);
-    const auto fileObject = file.ownerAs<QmlFile>();
-    if (!fileObject || !(fileObject && fileObject->isValid())) {
-        guard.setError({
-                int(QLspSpecification::ErrorCodes::RequestCancelled),
-                "Cannot proceed: current QML document is invalid!"_L1,
-        });
-        return;
-    }
-    auto newEncoded = QmlHighlighting::Utils::collectTokens(file, std::nullopt, m_mode);
-    auto &registeredTokens = m_codeModelManager->registeredTokens(uri);
-    const auto lastResultId = registeredTokens.resultId;
-    QmlHighlighting::Utils::updateResultID(registeredTokens.resultId);
-
-    // Return full token list if result ids not align
-    // otherwise compute the delta.
-    if (lastResultId == request->m_parameters.previousResultId) {
-        result = QLspSpecification::SemanticTokensDelta{
-            registeredTokens.resultId,
-            QmlHighlighting::Utils::computeDiff(registeredTokens.lastTokens, newEncoded)
-        };
-    } else if (!newEncoded.isEmpty()) {
-        result = QLspSpecification::SemanticTokens{ registeredTokens.resultId, newEncoded };
+    auto &cached = m_codeModelManager->registeredTokens(uri);
+    if (cached.resultId != request->m_parameters.previousResultId) {
+        // The client is out of sync, send full tokens
+        cached.resultId = request->m_parameters.previousResultId;
+        const auto encoded = generateHighlights(cached, doc, std::nullopt, m_mode);
+        result = QLspSpecification::SemanticTokens{ cached.resultId, encoded };
     } else {
-        result = nullptr;
+        const auto cachedHighlights = QmlHighlighting::Utils::encodeSemanticTokens(cached.highlights);
+        const auto encoded = generateHighlights(cached, doc, std::nullopt, m_mode);
+        result = QLspSpecification::SemanticTokensDelta{
+                cached.resultId,
+                QmlHighlighting::Utils::computeDiff(cachedHighlights, encoded)
+            };
     }
-    registeredTokens.lastTokens = std::move(newEncoded);
 }
 
 void SemanticTokenDeltaHandler::registerHandlers(QLanguageServer *, QLanguageServerProtocol *protocol)
@@ -153,7 +162,7 @@ https://microsoft.github.io/language-server-protocol/specifications/specificatio
 Sends a QLspSpecification::SemanticTokens data as response that is generated for a range of file.
 */
 SemanticTokenRangeHandler::SemanticTokenRangeHandler(QmlLsp::QQmlCodeModelManager *codeModelManager)
-    : QQmlBaseModule(codeModelManager), m_mode(QmlHighlighting::HighlightingMode::Default)
+    : QQmlBaseModule(codeModelManager), m_mode(HighlightingMode::Default)
 {
 }
 
@@ -169,28 +178,21 @@ void SemanticTokenRangeHandler::process(
     ResponseScopeGuard guard(result, request->m_response);
     const QByteArray uri = QQmlLSUtils::lspUriToQmlUrl(request->m_parameters.textDocument.uri);
     const auto doc = m_codeModelManager->openDocumentByUrl(uri);
-    DomItem file = doc.snapshot.doc.fileObject(GoTo::MostLikely);
-    const auto qmlFile = file.as<QmlFile>();
-    if (!qmlFile || !(qmlFile && qmlFile->isValid())) {
-        guard.setError({
-                int(QLspSpecification::ErrorCodes::RequestCancelled),
-                "Cannot proceed: current QML document is invalid!"_L1,
-        });
-        return;
-    }
-    const QString &code = qmlFile->code();
+    const QString code = doc.textDocument->toPlainText();
     const auto range = request->m_parameters.range;
     int startOffset =
             int(QQmlLSUtils::textOffsetFrom(code, range.start.line, range.end.character));
     int endOffset = int(QQmlLSUtils::textOffsetFrom(code, range.end.line, range.end.character));
-    auto &&encoded = QmlHighlighting::Utils::collectTokens(
-            file, QmlHighlighting::HighlightsRange{ startOffset, endOffset }, m_mode);
-    auto &registeredTokens = m_codeModelManager->registeredTokens(uri);
-    if (!encoded.isEmpty()) {
-        QmlHighlighting::Utils::updateResultID(registeredTokens.resultId);
-        result = SemanticTokens{ registeredTokens.resultId, std::move(encoded) };
-    } else {
+    auto &cached = m_codeModelManager->registeredTokens(uri);
+    const auto encodedTokens = generateHighlights(
+            cached,
+            doc,
+            QmlHighlighting::HighlightsRange{ startOffset, endOffset },
+            m_mode);
+    if (encodedTokens.isEmpty()) {
         result = nullptr;
+    } else {
+        result = SemanticTokens{ cached.resultId, std::move(encodedTokens) };
     }
 }
 
@@ -226,7 +228,7 @@ void QQmlHighlightSupport::setupCapabilities(
 
     if (auto clientInitOptions = clientCapabilities.initializationOptions) {
         if ((*clientInitOptions)[u"qtCreatorHighlighting"_s].toBool(false)) {
-            const auto mode = QmlHighlighting::HighlightingMode::QtCHighlighting;
+            const auto mode = HighlightingMode::QtCHighlighting;
             m_delta.setHighlightingMode(mode);
             m_full.setHighlightingMode(mode);
             m_range.setHighlightingMode(mode);
diff --git a/src/qmlls/qqmlsemantictokens.cpp b/src/qmlls/qqmlsemantictokens.cpp
index b5b38f827f..3a0157f447 100644
--- a/src/qmlls/qqmlsemantictokens.cpp
+++ b/src/qmlls/qqmlsemantictokens.cpp
@@ -1057,11 +1057,15 @@ HighlightsContainer Utils::visitTokens(const QQmlJS::Dom::DomItem &item,
     return highlightDomElements.highlights();
 }
 
-QList<int> Utils::collectTokens(const QQmlJS::Dom::DomItem &item,
-                                     const std::optional<HighlightsRange> &range,
-                                     HighlightingMode mode)
+HighlightsContainer Utils::shiftHighlights(const HighlightsContainer &cachedHighlights,
+                                           const QString &lastValidCode, const QString &currentCode)
 {
-    return Utils::encodeSemanticTokens(visitTokens(item, range), mode);
+    using namespace QQmlLSUtils;
+    Differ differ;
+    const QList<Diff> diffs = differ.diff(lastValidCode, currentCode);
+    HighlightsContainer shifts = cachedHighlights;
+    applyDiffs(shifts, diffs);
+    return shifts;
 }
 
 static std::pair<quint32, quint32> newlineCountAndLastLineLength(const QString &text)
@@ -1069,6 +1073,7 @@ static std::pair<quint32, quint32> newlineCountAndLastLineLength(const QString &
     auto [row, col]  = QQmlJS::SourceLocation::rowAndColumnFrom(text, text.size());
     return { row - 1, col - 1 }; // rows are 1-based, so subtract 1 to get the number of newlines
 }
+
 static void updateCursorPositionByDiff(const QString &text, QQmlJS::SourceLocation &cursor)
 {
     auto [newLines, lastLineLength] = newlineCountAndLastLineLength(text);
diff --git a/src/qmlls/qqmlsemantictokens_p.h b/src/qmlls/qqmlsemantictokens_p.h
index 1641c9f303..0674a0a3cc 100644
--- a/src/qmlls/qqmlsemantictokens_p.h
+++ b/src/qmlls/qqmlsemantictokens_p.h
@@ -160,14 +160,13 @@ void addModifier(QLspSpecification::SemanticTokenModifiers modifier, int *baseMo
 bool rangeOverlapsWithSourceLocation(const QQmlJS::SourceLocation &loc, const HighlightsRange &r);
 QList<QLspSpecification::SemanticTokensEdit> computeDiff(const QList<int> &, const QList<int> &);
 void updateResultID(QByteArray &resultID);
-QList<int> collectTokens(const QQmlJS::Dom::DomItem &item,
-                            const std::optional<HighlightsRange> &range,
-                            HighlightingMode mode = HighlightingMode::Default);
 HighlightsContainer visitTokens(const QQmlJS::Dom::DomItem &item,
                                 const std::optional<HighlightsRange> &range);
 void addHighlight(HighlightsContainer &out, const QQmlJS::SourceLocation &loc, QmlHighlightKind,
                     QmlHighlightModifiers = QmlHighlightModifier::None);
 void applyDiffs(HighlightsContainer &highlights, const QList<QQmlLSUtils::Diff> &diffs);
+HighlightsContainer shiftHighlights(const HighlightsContainer &cachedHighlights,
+                                    const QString &lastValidCode, const QString &currentCode);
 } // namespace Utils
 
 class HighlightingVisitor