diff --git a/src/compile/directive.cpp b/src/compile/directive.cpp
index bf9930964..dd6f45f1c 100644
--- a/src/compile/directive.cpp
+++ b/src/compile/directive.cpp
@@ -141,18 +141,6 @@ class DirectiveCollector : public clang::PPCallbacks {
         }
     }
 
-    void moduleImport(clang::SourceLocation import_location,
-                      clang::ModuleIdPath names,
-                      const clang::Module*) override {
-        auto fid = unit.file_id(unit.expansion_location(import_location));
-        auto& import = unit->directives[fid].imports.emplace_back();
-        import.location = import_location;
-        for(auto name: names) {
-            import.name += name.getIdentifierInfo()->getName();
-            import.name_locations.emplace_back(name.getLoc());
-        }
-    }
-
     void HasInclude(clang::SourceLocation location,
                     llvm::StringRef,
                     bool,
diff --git a/src/compile/directive.h b/src/compile/directive.h
index 34dc15c5d..6c0ddadc4 100644
--- a/src/compile/directive.h
+++ b/src/compile/directive.h
@@ -117,18 +117,6 @@ struct Pragma {
     clang::SourceLocation loc;
 };
 
-struct Import {
-    /// The name of imported module.
-    std::string name;
-
-    /// The location of import keyword, may comes from macro expansion.
-    clang::SourceLocation location;
-
-    /// The locations of tokens that make up the token name, may comes
-    /// from macro expansion.
-    std::vector<clang::SourceLocation> name_locations;
-};
-
 /// Information about `#embed` directive.
 struct Embed {
     /// The file name in the embed directive, not including quotes or angle brackets.
@@ -168,7 +156,6 @@ struct Directive {
     std::vector<Condition> conditions;
     std::vector<MacroRef> macros;
     std::vector<Pragma> pragmas;
-    std::vector<Import> imports;
     std::vector<Embed> embeds;
     std::vector<HasEmbed> has_embeds;
 };
diff --git a/src/feature/semantic_tokens.cpp b/src/feature/semantic_tokens.cpp
index 4f81ae82d..fbddc3a64 100644
--- a/src/feature/semantic_tokens.cpp
+++ b/src/feature/semantic_tokens.cpp
@@ -242,7 +242,13 @@ class SemanticTokensCollector : public SemanticVisitor<SemanticTokensCollector>
         add_token(location, SymbolKind::Macro, modifiers);
     }
 
-    // handleModuleOccurrence
+    void handleModuleOccurrence(clang::SourceLocation keyword,
+                                llvm::ArrayRef<clang::SourceLocation> identifiers) {
+        add_token(keyword, SymbolKind::Keyword, 0);
+        for(auto loc: identifiers) {
+            add_token(loc, SymbolKind::Module, 0);
+        }
+    }
 
     // handleRelation
 
@@ -296,16 +302,6 @@ class SemanticTokensCollector : public SemanticVisitor<SemanticTokensCollector>
     void highlight_modules() {
         auto interested = unit.interested_file();
 
-        auto directives_it = unit.directives().find(interested);
-        if(directives_it != unit.directives().end()) {
-            for(const auto& import: directives_it->second.imports) {
-                add_token(import.location, SymbolKind::Keyword, 0);
-                for(auto loc: import.name_locations) {
-                    add_token(loc, SymbolKind::Module, 0);
-                }
-            }
-        }
-
         auto* mod = unit.context().getCurrentNamedModule();
         if(!mod) {
             return;
diff --git a/src/semantic/semantic_visitor.h b/src/semantic/semantic_visitor.h
index 747291958..78fa6ff1a 100644
--- a/src/semantic/semantic_visitor.h
+++ b/src/semantic/semantic_visitor.h
@@ -63,13 +63,11 @@ class SemanticVisitor : public FilteredASTVisitor<SemanticVisitor<Derived>> {
 
     /// Invoked when a module occurrence is seen in source code.
     /// @param keyword The location of the `module` or `import` keyword.
-    /// @param identifiers Tokens that make up the module name.
+    /// @param identifiers Source locations of identifiers that make up the module name.
     void handleModuleOccurrence(clang::SourceLocation keyword,
-                                llvm::ArrayRef<clang::syntax::Token> identifiers) {
+                                llvm::ArrayRef<clang::SourceLocation> identifiers) {
         assert(keyword.isValid() && keyword.isFileID() && "Invalid keyword location");
 
-        /// FIXME: Check whether identifiers are valid.
-
         if constexpr(!std::same_as<decltype(&SemanticVisitor::handleModuleOccurrence),
                                    decltype(&Derived::handleModuleOccurrence)>) {
             getDerived().handleModuleOccurrence(keyword, identifiers);
@@ -145,30 +143,16 @@ class SemanticVisitor : public FilteredASTVisitor<SemanticVisitor<Derived>> {
 #define VISIT_TYPELOC(type) bool Visit##type(clang::type loc)
 
     VISIT_DECL(ImportDecl) {
-        /// FIXME:
-        // auto tokens = TB.expandedTokens(decl->getSourceRange());
-        //
-        // assert(tokens.size() >= 2 && tokens[0].kind() == clang::tok::identifier &&
-        //       tokens[0].text(SM) == "import" && "Invalid import declaration");
-        // assert([&]() {
-        //    auto range = tokens.drop_front(1);
-        //    for(auto iter = range.begin(); iter != range.end(); ++iter) {
-        //        if(iter->kind() == clang::tok::identifier) {
-        //            if(auto next = iter + 1;
-        //               next != range.end() && (next->kind() == clang::tok::coloncolon ||
-        //                                       next->kind() == clang::tok::period)) {
-        //                continue;
-        //            }
-        //            break;
-        //        } else {
-        //            return false;
-        //        }
-        //    }
-        //    return true;
-        //}() && "Invalid import declaration");
-        //
-        // handleModuleOccurrence(tokens[0].location(), tokens.drop_front(1));
+        auto keyword = decl->getLocation();
+        auto tokens = unit.expanded_tokens(decl->getSourceRange());
+        for(const auto& token: tokens) {
+            if(token.text(unit.context().getSourceManager()) == "import") {
+                keyword = token.location();
+                break;
+            }
+        }
 
+        handleModuleOccurrence(keyword, decl->getIdentifierLocs());
         return true;
     }
 
diff --git a/tests/unit/feature/semantic_tokens_tests.cpp b/tests/unit/feature/semantic_tokens_tests.cpp
index e83aac897..f1f7d1182 100644
--- a/tests/unit/feature/semantic_tokens_tests.cpp
+++ b/tests/unit/feature/semantic_tokens_tests.cpp
@@ -496,6 +496,24 @@ export @kw[import] @mod[foo];
     EXPECT_TOKEN("mod", SymbolKind::Module);
 }
 
+TEST_CASE(ModulePartitionImport) {
+    add_files("main.cppm", R"(
+#[part.cppm]
+export module foo:part;
+export int x = 42;
+
+#[main.cppm]
+export module foo;
+@kw[import] :@mod[part];
+)");
+    ASSERT_TRUE(compile_with_modules());
+    tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF8);
+    decoded = decode_utf8_tokens(unit->interested_content(), tokens);
+
+    EXPECT_TOKEN("kw", SymbolKind::Keyword);
+    EXPECT_TOKEN("mod", SymbolKind::Module);
+}
+
 TEST_CASE(GlobalModuleFragment) {
     add_main("main.cpp", R"cpp(
 module;
diff --git a/tests/unit/test/tester.cpp b/tests/unit/test/tester.cpp
index e5b960427..a40f3806c 100644
--- a/tests/unit/test/tester.cpp
+++ b/tests/unit/test/tester.cpp
@@ -218,10 +218,16 @@ bool Tester::compile_with_modules(llvm::StringRef standard) {
         builder.params.vfs = overlay;
         builder.params.pcms = built_pcms;
 
-        if(!builder.try_compile())
+        PCMInfo info;
+        auto built = clice::compile(builder.params, info);
+        if(!built.completed()) {
+            for(auto& diag: built.diagnostics()) {
+                LOG_ERROR("{}", diag.message);
+            }
             return false;
+        }
 
-        built_pcms.try_emplace(mod.module_name, *pcm_path);
+        built_pcms.try_emplace(mod.module_name, info.path);
     }
 
     prepare(standard);