feat: import auto-completion (#2904)

This PR adds basic auto-completion support for imports. Since it still lacks Lake support for accurate completion suggestions (cc @tydeu - we already know what needs to be done), it falls back to traversing the `LEAN_SRC_PATH` for available imports. Three kinds of import completion requests are supported: - Completion of the full `import` command. Triggered when requesting completions in an empty space within the header. - Known issue: It is possible to trigger this completion within a comment in the header. Fixing this would require architecture for parsing some kind of sub-syntax between individual commands. - Completion of the full module name after an incomplete `import` command. - Completion of a partial module name with a trailing dot. Since the set of imports is potentially expensive to compute, they are cached for 10 seconds after the last import auto-completion request. Closes #2655. ### Changes This PR also makes the following changes: - To support completions on the trailing dot, the `import` syntax was adjusted to provide partial syntax when a trailing dot is used. - `FileWorker.lean` was refactored lightly with some larger definitions being broken apart. - The `WorkerState` gained two new fields: - `currHeaderStx` tracks the current header syntax, as opposed to tracking only the initial header syntax in `initHeaderStx`. When the header syntax changes, a task is launched that restarts the file worker after a certain delay to avoid constant restarts while editing the header. During this time period, we may still want to serve import auto-completion requests, so we need to know the up-to-date header syntax. - `importCachingTask?` contains a task that computes the set of available imports. - `determineLakePath` has moved to a new file `Lean/Util/LakePath.lean` as it is now needed both in `ImportCompletion.lean` and `FileWorker.lean`. - `forEachModuleIn` from `Lake/Config/Blob.lean` has moved to `Lean/Util/Path.lean` as it is a generally useful utility function that was useful for traversing the `LEAN_SRC_PATH` as well. ### Tests Unfortunately, this PR lacks tests since the set of imports available in `tests/lean/interactive` will not be stable. In the future, I will add support for testing LSP requests in full project setups, which is when tests for import auto-completion will be added as well.
2023-11-24 08:46:19 +01:00 · 2023-11-24 08:46:19 +01:00 · 681fca1f8f
commit 681fca1f8f
parent e34656ce75
9 changed files with 402 additions and 146 deletions
--- a/src/Lean/Data/NameTrie.lean
+++ b/src/Lean/Data/NameTrie.lean
@ -69,4 +69,10 @@ def NameTrie.forMatchingM [Monad m] (t : NameTrie β) (k : Name) (f : β → m U
 def NameTrie.forM [Monad m] (t : NameTrie β) (f : β → m Unit) : m Unit :=
  t.forMatchingM Name.anonymous f

+def NameTrie.matchingToArray (t : NameTrie β) (k : Name) : Array β :=
+  Id.run <| t.foldMatchingM k #[] fun v acc => acc.push v
+
+def NameTrie.toArray (t : NameTrie β) : Array β :=
+  Id.run <| t.foldM #[] fun v acc => acc.push v
+
 end Lean
--- a/src/Lean/Parser/Module.lean
+++ b/src/Lean/Parser/Module.lean
@ -11,7 +11,11 @@ namespace Parser

 namespace Module
 def «prelude»  := leading_parser "prelude"
-def «import»   := leading_parser "import " >> optional "runtime" >> ident
+-- `optional (checkNoWsBefore >> "." >> checkNoWsBefore >> ident)`
+-- can never fully succeed but ensures that `import (runtime)? <ident>.`
+-- produces a partial syntax that contains the dot.
+-- The partial syntax is useful for import dot-auto-completion.
+def «import»   := leading_parser "import " >> optional "runtime" >> ident >> optional (checkNoWsBefore >> "." >> checkNoWsBefore >> ident)
 def header     := leading_parser optional («prelude» >> ppLine) >> many («import» >> ppLine) >> ppLine
 /--
  Parser for a Lean module. We never actually run this parser but instead use the imperative definitions below that
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@ -6,7 +6,6 @@ Authors: Marc Huisinga, Wojciech Nawrocki
 -/
 import Init.System.IO
 import Lean.Data.RBMap
-
 import Lean.Environment

 import Lean.Data.Lsp
@ -25,6 +24,7 @@ import Lean.Server.FileWorker.RequestHandling
 import Lean.Server.FileWorker.WidgetRequests
 import Lean.Server.Rpc.Basic
 import Lean.Widget.InteractiveDiagnostic
+import Lean.Server.ImportCompletion

 /-!
 For general server architecture, see `README.md`. For details of IPC communication, see `Watchdog.lean`.
@ -125,7 +125,7 @@ section Elab
    let ctx ← read
    let some headerSnap := snaps[0]? | panic! "empty snapshots"
    if headerSnap.msgLog.hasErrors then
-      -- Treat header processing errors as fatal so users aren't swamped with
+      -- Treats header processing errors as fatal so users aren't swamped with
      -- followup errors
      publishProgressAtPos m headerSnap.beginPos ctx.hOut (kind := LeanFileProgressKind.fatalError)
      publishIleanInfoFinal m ctx.hOut #[headerSnap]
@ -142,35 +142,47 @@ end Elab
 -- Pending requests are tracked so they can be cancelled
 abbrev PendingRequestMap := RBMap RequestID (Task (Except IO.Error Unit)) compare

+structure AvailableImportsCache where
+  availableImports       : ImportCompletion.AvailableImports
+  lastRequestTimestampMs : Nat
+
 structure WorkerState where
-  doc             : EditableDocument
-  -- The initial header syntax tree. Changing the header requires restarting the worker process.
-  initHeaderStx   : Syntax
-  pendingRequests : PendingRequestMap
+  doc                : EditableDocument
+  -- The initial header syntax tree that the file worker was started with.
+  initHeaderStx      : Syntax
+  -- The current header syntax tree. Changing the header from `initHeaderStx` initiates a restart
+  -- that only completes after a while, so `currHeaderStx` tracks the modified syntax until then.
+  currHeaderStx      : Syntax
+  importCachingTask? : Option (Task (Except Error AvailableImportsCache))
+  pendingRequests    : PendingRequestMap
  /-- A map of RPC session IDs. We allow asynchronous elab tasks and request handlers
  to modify sessions. A single `Ref` ensures atomic transactions. -/
-  rpcSessions     : RBMap UInt64 (IO.Ref RpcSession) compare
+  rpcSessions        : RBMap UInt64 (IO.Ref RpcSession) compare

 abbrev WorkerM := ReaderT WorkerContext <| StateRefT WorkerState IO

 /- Worker initialization sequence. -/
 section Initialization
-  /-- Use `lake print-paths` to compile dependencies on the fly and add them to `LEAN_PATH`.
-  Compilation progress is reported to `hOut` via LSP notifications. Return the search path for
-  source files. -/
-  partial def lakeSetupSearchPath (lakePath : System.FilePath) (m : DocumentMeta) (imports : Array Import) (hOut : FS.Stream) : IO SearchPath := do
+
+  structure LakePrintPathsResult where
+    spawnArgs : Process.SpawnArgs
+    exitCode  : UInt32
+    stdout    : String
+    stderr    : String
+
+  partial def runLakePrintPaths (lakePath : System.FilePath) (m : DocumentMeta) (imports : Array Import) (hOut : FS.Stream) : IO LakePrintPathsResult := do
    let mut args := #["print-paths"] ++ imports.map (toString ·.module)
    if m.dependencyBuildMode matches .never then
      args := args.push "--no-build"
-    let cmdStr := " ".intercalate (toString lakePath :: args.toList)
-    let lakeProc ← Process.spawn {
+    let spawnArgs : Process.SpawnArgs := {
      stdin  := Process.Stdio.null
      stdout := Process.Stdio.piped
      stderr := Process.Stdio.piped
      cmd    := lakePath.toString
      args
    }
-    -- progress notification: report latest stderr line
+    let lakeProc ← Process.spawn spawnArgs
+    -- progress notification: reports latest stderr line
    let rec processStderr (acc : String) : IO String := do
      let line ← lakeProc.stderr.getLine
      if line == "" then
@ -181,65 +193,97 @@ section Initialization
    let stderr ← IO.asTask (processStderr "") Task.Priority.dedicated
    let stdout := String.trim (← lakeProc.stdout.readToEnd)
    let stderr ← IO.ofExcept stderr.get
-    match (← lakeProc.wait) with
+    let exitCode ← lakeProc.wait
+    return ⟨spawnArgs, exitCode, stdout, stderr⟩
+
+  /-- Uses `lake print-paths` to compile dependencies on the fly and adds them to `LEAN_PATH`.
+  Compilation progress is reported to `hOut` via LSP notifications. Returns the search path for
+  source files. -/
+  partial def lakeSetupSearchPath (lakePath : System.FilePath) (m : DocumentMeta) (imports : Array Import) (hOut : FS.Stream) : IO SearchPath := do
+    let result ← runLakePrintPaths lakePath m imports hOut
+    let cmdStr := " ".intercalate (toString result.spawnArgs.cmd :: result.spawnArgs.args.toList)
+    match result.exitCode with
    | 0 =>
-      let Except.ok (paths : LeanPaths) ← pure (Json.parse stdout >>= fromJson?)
-        | throwServerError s!"invalid output from `{cmdStr}`:\n{stdout}\nstderr:\n{stderr}"
+      let Except.ok (paths : LeanPaths) ← pure (Json.parse result.stdout >>= fromJson?)
+        | throwServerError s!"invalid output from `{cmdStr}`:\n{result.stdout}\nstderr:\n{result.stderr}"
      initSearchPath (← getBuildDir) paths.oleanPath
      paths.loadDynlibPaths.forM loadDynlib
      paths.srcPath.mapM realPathNormalized
    | 2 => pure []  -- no lakefile.lean
    -- error from `--no-build`
-    | 3 => throwServerError s!"Imports are out of date and must be rebuilt; use the \"Restart File\" command in your editor.\n\n{stdout}"
-    | _ => throwServerError s!"`{cmdStr}` failed:\n{stdout}\nstderr:\n{stderr}"
+    | 3 => throwServerError s!"Imports are out of date and must be rebuilt; use the \"Restart File\" command in your editor.\n\n{result.stdout}"
+    | _ => throwServerError s!"`{cmdStr}` failed:\n{result.stdout}\nstderr:\n{result.stderr}"
+
+  def setupSrcSearchPath (m : DocumentMeta) (hOut : FS.Stream) (headerStx : Syntax) : IO SearchPath := do
+    let some path := System.Uri.fileUriToPath? m.uri
+      | return ← initSrcSearchPath
+
+    -- NOTE: we assume for now that `lakefile.lean` does not have any non-core-Lean deps
+    -- NOTE: lake does not exist in stage 0 (yet?)
+    if path.fileName == "lakefile.lean" then
+      return ← initSrcSearchPath
+
+    let lakePath ← determineLakePath
+    if !(← System.FilePath.pathExists lakePath) then
+      return ← initSrcSearchPath
+
+    let imports := Lean.Elab.headerToImports headerStx
+    let pkgSearchPath ← lakeSetupSearchPath lakePath m imports hOut
+    return ← initSrcSearchPath pkgSearchPath
+
+  def buildHeaderEnv (m : DocumentMeta) (hOut : FS.Stream) (opts : Options) (headerStx : Syntax) : IO (Environment × MessageLog × SearchPath) := do
+    let (headerEnv, msgLog, srcSearchPath) ←
+      match ← EIO.toIO' <| setupSrcSearchPath m hOut headerStx with
+      | .ok srcSearchPath =>
+        -- allows `headerEnv` to be leaked, which would live until the end of the process anyway
+        let (headerEnv, msgLog) ← Elab.processHeader (leakEnv := true) headerStx opts MessageLog.empty m.mkInputContext
+        pure (headerEnv, msgLog, srcSearchPath)
+      | .error e =>
+        let msgs := MessageLog.empty.add { fileName := "<ignored>", pos := ⟨0, 0⟩, data := e.toString }
+        pure (← mkEmptyEnvironment, msgs, ← initSrcSearchPath)
+
+    let mut headerEnv := headerEnv
+    try
+      if let some path := System.Uri.fileUriToPath? m.uri then
+        headerEnv := headerEnv.setMainModule (← moduleNameOfFileName path none)
+    catch _ =>
+      pure ()
+
+    return (headerEnv, msgLog, srcSearchPath)
+
+  def buildCommandState
+      (m : DocumentMeta)
+      (headerStx : Syntax)
+      (headerEnv : Environment)
+      (headerMsgLog : MessageLog)
+      (opts : Options)
+      : Elab.Command.State :=
+    let headerContextInfo : Elab.ContextInfo := {
+      env     := headerEnv
+      fileMap := m.text
+      ngen    := { namePrefix := `_worker }
+    }
+    let headerInfo := Elab.Info.ofCommandInfo { elaborator := `header, stx := headerStx }
+    let headerInfoNodes := headerStx[1].getArgs.toList.map fun importStx =>
+      Elab.InfoTree.node (Elab.Info.ofCommandInfo {
+        elaborator := `import
+        stx := importStx
+      }) #[].toPArray'
+    let headerInfoTree := Elab.InfoTree.node headerInfo headerInfoNodes.toPArray'
+    let headerInfoState := {
+      enabled := true
+      trees := #[Elab.InfoTree.context headerContextInfo headerInfoTree].toPArray'
+    }
+    { Elab.Command.mkState headerEnv headerMsgLog opts with infoState := headerInfoState }

  def compileHeader (m : DocumentMeta) (hOut : FS.Stream) (opts : Options) (hasWidgets : Bool)
      : IO (Syntax × Task (Except Error (Snapshot × SearchPath))) := do
    -- parsing should not take long, do synchronously
-    let (headerStx, headerParserState, msgLog) ← Parser.parseHeader m.mkInputContext
+    let (headerStx, headerParserState, parseMsgLog) ← Parser.parseHeader m.mkInputContext
    (headerStx, ·) <$> EIO.asTask do
-      let mut srcSearchPath ← initSrcSearchPath (← getBuildDir)
-      let lakePath ← match (← IO.getEnv "LAKE") with
-        | some path => pure <| System.FilePath.mk path
-        | none =>
-          let lakePath ← match (← IO.getEnv "LEAN_SYSROOT") with
-            | some path => pure <| System.FilePath.mk path / "bin" / "lake"
-            | _         => pure <| (← appDir) / "lake"
-          pure <| lakePath.withExtension System.FilePath.exeExtension
-      let (headerEnv, msgLog) ← try
-        if let some path := System.Uri.fileUriToPath? m.uri then
-          -- NOTE: we assume for now that `lakefile.lean` does not have any non-stdlib deps
-          -- NOTE: lake does not exist in stage 0 (yet?)
-          if path.fileName != "lakefile.lean" && (← System.FilePath.pathExists lakePath) then
-            let pkgSearchPath ← lakeSetupSearchPath lakePath m (Lean.Elab.headerToImports headerStx) hOut
-            srcSearchPath ← initSrcSearchPath (← getBuildDir) pkgSearchPath
-        -- allow `headerEnv` to be leaked, which would live until the end of the process anyway
-        Elab.processHeader (leakEnv := true) headerStx opts msgLog m.mkInputContext
-      catch e =>  -- should be from `lake print-paths`
-        let msgs := MessageLog.empty.add { fileName := "<ignored>", pos := ⟨0, 0⟩, data := e.toString }
-        pure (← mkEmptyEnvironment, msgs)
-      let mut headerEnv := headerEnv
-      try
-        if let some path := System.Uri.fileUriToPath? m.uri then
-          headerEnv := headerEnv.setMainModule (← moduleNameOfFileName path none)
-      catch _ => pure ()
-      let cmdState := Elab.Command.mkState headerEnv msgLog opts
-      let cmdState := { cmdState with infoState := {
-        enabled := true
-        trees := #[Elab.InfoTree.context ({
-          env     := headerEnv
-          fileMap := m.text
-          ngen    := { namePrefix := `_worker }
-        }) (Elab.InfoTree.node
-            (Elab.Info.ofCommandInfo { elaborator := `header, stx := headerStx })
-            (headerStx[1].getArgs.toList.map (fun importStx =>
-              Elab.InfoTree.node (Elab.Info.ofCommandInfo {
-                elaborator := `import
-                stx := importStx
-              }) #[].toPArray'
-            )).toPArray'
-        )].toPArray'
-      }}
+      let (headerEnv, envMsgLog, srcSearchPath) ← buildHeaderEnv m hOut opts headerStx
+      let headerMsgLog := parseMsgLog.append envMsgLog
+      let cmdState := buildCommandState m headerStx headerEnv headerMsgLog opts
      let headerSnap := {
        beginPos := 0
        stx := headerStx
@ -256,23 +300,25 @@ section Initialization
    let clientHasWidgets := initParams.initializationOptions?.bind (·.hasWidgets?) |>.getD false
    let (headerStx, headerTask) ← compileHeader meta o opts (hasWidgets := clientHasWidgets)
    let cancelTk ← CancelToken.new
-    let ctx :=
-      { hIn  := i
-        hOut := o
-        hLog := e
-        headerTask
-        initParams
-        clientHasWidgets
-      }
+    let ctx := {
+      hIn  := i
+      hOut := o
+      hLog := e
+      headerTask
+      initParams
+      clientHasWidgets
+    }
    let cmdSnaps ← EIO.mapTask (t := headerTask) (match · with
      | Except.ok (s, _) => unfoldCmdSnaps meta #[s] cancelTk ctx (startAfterMs := 0)
      | Except.error e   => throw (e : ElabTaskError))
    let doc : EditableDocument := { meta, cmdSnaps := AsyncList.delayed cmdSnaps, cancelTk }
-    return (ctx,
-    { doc             := doc
-      initHeaderStx   := headerStx
-      pendingRequests := RBMap.empty
-      rpcSessions     := RBMap.empty
+    return (ctx, {
+      doc                := doc
+      initHeaderStx      := headerStx
+      currHeaderStx      := headerStx
+      importCachingTask? := none
+      pendingRequests    := RBMap.empty
+      rpcSessions        := RBMap.empty
    })
 end Initialization

@ -280,53 +326,67 @@ section Updates
  def updatePendingRequests (map : PendingRequestMap → PendingRequestMap) : WorkerM Unit := do
    modify fun st => { st with pendingRequests := map st.pendingRequests }

-  /-- Given the new document, updates editable doc state. -/
-  def updateDocument (newMeta : DocumentMeta) : WorkerM Unit := do
+  def determineValidSnapshots (oldDoc : EditableDocument) (newMeta : DocumentMeta) (newHeaderSnap : Snapshot) : IO (List Snapshot) := do
+    let changePos := oldDoc.meta.text.source.firstDiffPos newMeta.text.source
+    -- Ignores exceptions, we are only interested in the successful snapshots
+    let (cmdSnaps, _) ← oldDoc.cmdSnaps.getFinishedPrefix
+    oldDoc.cmdSnaps.cancel
+    -- NOTE(WN): we invalidate eagerly as `endPos` consumes input greedily. To re-elaborate only
+    -- when really necessary, we could do a whitespace-aware `Syntax` comparison instead.
+    let mut validSnaps ← pure (cmdSnaps.takeWhile (fun s => s.endPos < changePos))
+    if h : validSnaps.length ≤ 1 then
+      validSnaps := [newHeaderSnap]
+    else
+      /- When at least one valid non-header snap exists, it may happen that a change does not fall
+          within the syntactic range of that last snap but still modifies it by appending tokens.
+          We check for this here. We do not currently handle crazy grammars in which an appended
+          token can merge two or more previous commands into one. To do so would require reparsing
+          the entire file. -/
+      have : validSnaps.length ≥ 2 := Nat.gt_of_not_le h
+      let mut lastSnap := validSnaps.getLast (by subst ·; simp at h)
+      let preLastSnap :=
+        have : 0 < validSnaps.length := Nat.lt_of_lt_of_le (by decide) this
+        have : validSnaps.length - 2 < validSnaps.length := Nat.sub_lt this (by decide)
+        validSnaps[validSnaps.length - 2]
+      let newLastStx ← parseNextCmd newMeta.mkInputContext preLastSnap
+      if newLastStx != lastSnap.stx then
+        validSnaps := validSnaps.dropLast
+    return validSnaps
+
+  def startNewSnapshotTasks (newMeta : DocumentMeta) : WorkerM (AsyncList ElabTaskError Snapshot × CancelToken) := do
    let ctx ← read
-    let oldDoc := (←get).doc
+    let oldDoc := (← get).doc
    oldDoc.cancelTk.set
    let initHeaderStx := (← get).initHeaderStx
    let (newHeaderStx, newMpState, _) ← Parser.parseHeader newMeta.mkInputContext
    let cancelTk ← CancelToken.new
-    let headSnapTask := oldDoc.cmdSnaps.waitHead?
-    let newSnaps ← if initHeaderStx != newHeaderStx then
-      EIO.asTask (ε := ElabTaskError) (prio := .dedicated) do
+
+    if initHeaderStx != newHeaderStx then
+      set { ← get with currHeaderStx := newHeaderStx }
+      let terminationTask ← EIO.asTask (ε := ElabTaskError) (prio := .dedicated) do
        IO.sleep ctx.initParams.editDelay.toUInt32
        cancelTk.check
        IO.Process.exit 2
-    else EIO.mapTask (ε := ElabTaskError) (t := headSnapTask) (prio := .dedicated) fun headSnap?? => do
+      return (AsyncList.delayed terminationTask, cancelTk)
+
+    let headSnapTask := oldDoc.cmdSnaps.waitHead?
+    let newSnapTasks ← EIO.mapTask (ε := ElabTaskError) (t := headSnapTask) (prio := .dedicated) fun headSnap?? => do
      -- There is always at least one snapshot absent exceptions
      let some headSnap ← MonadExcept.ofExcept headSnap?? | panic! "empty snapshots"
      let newHeaderSnap := { headSnap with stx := newHeaderStx, mpState := newMpState }
-      let changePos := oldDoc.meta.text.source.firstDiffPos newMeta.text.source
-      -- Ignore exceptions, we are only interested in the successful snapshots
-      let (cmdSnaps, _) ← oldDoc.cmdSnaps.getFinishedPrefix
-      oldDoc.cmdSnaps.cancel
-      -- NOTE(WN): we invalidate eagerly as `endPos` consumes input greedily. To re-elaborate only
-      -- when really necessary, we could do a whitespace-aware `Syntax` comparison instead.
-      let mut validSnaps ← pure (cmdSnaps.takeWhile (fun s => s.endPos < changePos))
-      if h : validSnaps.length ≤ 1 then
-        validSnaps := [newHeaderSnap]
-      else
-        /- When at least one valid non-header snap exists, it may happen that a change does not fall
-           within the syntactic range of that last snap but still modifies it by appending tokens.
-           We check for this here. We do not currently handle crazy grammars in which an appended
-           token can merge two or more previous commands into one. To do so would require reparsing
-           the entire file. -/
-        have : validSnaps.length ≥ 2 := Nat.gt_of_not_le h
-        let mut lastSnap := validSnaps.getLast (by subst ·; simp at h)
-        let preLastSnap :=
-          have : 0 < validSnaps.length := Nat.lt_of_lt_of_le (by decide) this
-          have : validSnaps.length - 2 < validSnaps.length := Nat.sub_lt this (by decide)
-          validSnaps[validSnaps.length - 2]
-        let newLastStx ← parseNextCmd newMeta.mkInputContext preLastSnap
-        if newLastStx != lastSnap.stx then
-          validSnaps := validSnaps.dropLast
+      let validSnaps ← determineValidSnapshots oldDoc newMeta newHeaderSnap
      -- wait for a bit, giving the initial `cancelTk.check` in `nextCmdSnap` time to trigger
      -- before kicking off any expensive elaboration (TODO: make expensive elaboration cancelable)
      unfoldCmdSnaps newMeta validSnaps.toArray cancelTk ctx
        (startAfterMs := ctx.initParams.editDelay.toUInt32)
-    modify fun st => { st with doc := { meta := newMeta, cmdSnaps := AsyncList.delayed newSnaps, cancelTk } }
+
+    return (AsyncList.delayed newSnapTasks, cancelTk)
+
+  /-- Given the new document, updates editable doc state. -/
+  def updateDocument (newMeta : DocumentMeta) : WorkerM Unit := do
+    let (newSnaps, cancelTk) ← startNewSnapshotTasks newMeta
+    modify fun st => { st with doc := { meta := newMeta, cmdSnaps := newSnaps, cancelTk } }
+
 end Updates

 /- Notifications are handled in the main thread. They may change global worker state
@ -396,6 +456,30 @@ section MessageHandling
      : WorkerM Unit := do
    updatePendingRequests (fun pendingRequests => pendingRequests.insert id requestTask)

+  def handleImportCompletionRequest (id : RequestID) (params : CompletionParams)
+      : WorkerM (Task (Except Error AvailableImportsCache)) := do
+    let ctx ← read
+    let st ← get
+    let text := st.doc.meta.text
+
+    match st.importCachingTask? with
+    | none => IO.asTask do
+      let availableImports ← ImportCompletion.collectAvailableImports
+      let lastRequestTimestampMs ← IO.monoMsNow
+      let completions := ImportCompletion.find text st.currHeaderStx params availableImports
+      ctx.hOut.writeLspResponse ⟨id, completions⟩
+      pure { availableImports, lastRequestTimestampMs : AvailableImportsCache }
+
+    | some task => IO.mapTask (t := task) fun result => do
+      let mut ⟨availableImports, lastRequestTimestampMs⟩ ← IO.ofExcept result
+      let timestampNowMs ← IO.monoMsNow
+      if timestampNowMs - lastRequestTimestampMs >= 10000 then
+        availableImports ← ImportCompletion.collectAvailableImports
+      lastRequestTimestampMs := timestampNowMs
+      let completions := ImportCompletion.find text st.currHeaderStx params availableImports
+      ctx.hOut.writeLspResponse ⟨id, completions⟩
+      pure { availableImports, lastRequestTimestampMs : AvailableImportsCache }
+
  def handleRequest (id : RequestID) (method : String) (params : Json)
      : WorkerM Unit := do
    let ctx ← read
@ -413,26 +497,33 @@ section MessageHandling
            message := toString e }
      return

+    if method == "textDocument/completion" then
+      let params ← parseParams CompletionParams params
+      if ImportCompletion.isImportCompletionRequest st.doc.meta.text st.currHeaderStx params then
+        let importCachingTask ← handleImportCompletionRequest id params
+        set <| { st with importCachingTask? := some importCachingTask }
+        return
+
    -- we assume that every request requires at least the header snapshot or the search path
    let t ← IO.bindTask ctx.headerTask fun x => do
-     let (_, srcSearchPath) ← IO.ofExcept x
-     let rc : RequestContext :=
-       { rpcSessions := st.rpcSessions
-         srcSearchPath
-         doc := st.doc
-         hLog := ctx.hLog
-         hOut := ctx.hOut
-         initParams := ctx.initParams }
-     let t? ← EIO.toIO' <| handleLspRequest method params rc
-     let t₁ ← match t? with
-       | Except.error e =>
-         IO.asTask do
-           ctx.hOut.writeLspResponseError <| e.toLspResponseError id
-       | Except.ok t => (IO.mapTask · t) fun
-         | Except.ok resp =>
-           ctx.hOut.writeLspResponse ⟨id, resp⟩
-         | Except.error e =>
-           ctx.hOut.writeLspResponseError <| e.toLspResponseError id
+      let (_, srcSearchPath) ← IO.ofExcept x
+      let rc : RequestContext :=
+        { rpcSessions := st.rpcSessions
+          srcSearchPath
+          doc := st.doc
+          hLog := ctx.hLog
+          hOut := ctx.hOut
+          initParams := ctx.initParams }
+      let t? ← EIO.toIO' <| handleLspRequest method params rc
+      let t₁ ← match t? with
+        | Except.error e =>
+          IO.asTask do
+            ctx.hOut.writeLspResponseError <| e.toLspResponseError id
+        | Except.ok t => (IO.mapTask · t) fun
+          | Except.ok resp =>
+            ctx.hOut.writeLspResponse ⟨id, resp⟩
+          | Except.error e =>
+            ctx.hOut.writeLspResponseError <| e.toLspResponseError id
    queueRequest id t
 end MessageHandling

--- a/src/Lean/Server/ImportCompletion.lean
+++ b/src/Lean/Server/ImportCompletion.lean
@ -0,0 +1,141 @@
+/-
+Copyright (c) 2023 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Marc Huisinga
+-/
+import Lean.Data.Name
+import Lean.Data.NameTrie
+import Lean.Data.Lsp.Utf16
+import Lean.Data.Lsp.LanguageFeatures
+import Lean.Util.Paths
+import Lean.Util.LakePath
+
+namespace ImportCompletion
+
+open Lean Lsp
+
+abbrev ImportTrie := Lean.NameTrie Name
+
+abbrev AvailableImports := Array Name
+
+def AvailableImports.toImportTrie (imports : AvailableImports) : ImportTrie := Id.run do
+  let mut importTrie := ∅
+  for i in imports do
+    importTrie := importTrie.insert i i
+  return importTrie
+
+def determinePartialHeaderCompletions
+    (headerStx : Syntax)
+    (completionPos : String.Pos)
+    : Option Syntax := Id.run do
+  let some importCmdToComplete := headerStx[1].find? fun importStx => Id.run do
+      let importIdStx := importStx
+      let some startPos := importIdStx.getPos?
+        | return false
+      let some endPos := importIdStx.getTailPos?
+        | return false
+      return startPos <= completionPos && completionPos <= endPos
+    | return none
+  return some importCmdToComplete
+
+/-- Checks whether `completionPos` points at the position after an incomplete `import` statement. -/
+def isImportNameCompletionRequest (headerStx : Syntax) (completionPos : String.Pos) : Bool :=
+  headerStx[1].getArgs.any fun importStx =>
+    let importCmd := importStx[0]
+    let importId := importStx[2]
+    importId.isMissing && importCmd.getTailPos?.isSome && completionPos == importCmd.getTailPos?.get! + ' '
+
+/-- Checks whether `completionPos` points at a free space in the header. -/
+def isImportCmdCompletionRequest (headerStx : Syntax) (completionPos : String.Pos) : Bool :=
+  ! headerStx[1].getArgs.any fun importStx => importStx.getArgs.any fun arg =>
+    arg.getPos?.isSome && arg.getTailPos?.isSome
+      && arg.getPos?.get! <= completionPos && completionPos <= arg.getTailPos?.get!
+
+def computePartialImportCompletions
+    (headerStx : Syntax)
+    (completionPos : String.Pos)
+    (availableImports : ImportTrie)
+    : Array Name := Id.run do
+  let some importStxToComplete := headerStx[1].getArgs.find? fun importStx => Id.run do
+      -- `partialTrailingDotStx` ≙ `("." ident)?`
+      let partialTrailingDotStx := importStx[3]
+      if ! partialTrailingDotStx.hasArgs then
+        return false
+      let trailingDot := partialTrailingDotStx[0]
+      let some tailPos := trailingDot.getTailPos?
+        | return false
+      return tailPos == completionPos
+    | return #[]
+  let importPrefixToComplete := importStxToComplete[2].getId
+
+  let completions : Array Name :=
+    availableImports.matchingToArray importPrefixToComplete
+      |>.map fun matchingAvailableImport =>
+        matchingAvailableImport.replacePrefix importPrefixToComplete Name.anonymous
+
+  let nonEmptyCompletions := completions.filter fun completion => !completion.isAnonymous
+
+  return nonEmptyCompletions.insertionSort (Name.cmp · · == Ordering.lt)
+
+def isImportCompletionRequest (text : FileMap) (headerStx : Syntax) (params : CompletionParams) : Bool :=
+  let completionPos := text.lspPosToUtf8Pos params.position
+  let headerStartPos := headerStx.getPos?.getD 0
+  let headerEndPos := headerStx.getTailPos?.getD headerStartPos
+  completionPos <= headerEndPos + ' ' + ' '
+
+def collectAvailableImportsFromLake : IO (Option AvailableImports) := do
+  let lakePath ← Lean.determineLakePath
+  let spawnArgs : IO.Process.SpawnArgs := {
+    stdin  := IO.Process.Stdio.null
+    stdout := IO.Process.Stdio.piped
+    stderr := IO.Process.Stdio.piped
+    cmd    := lakePath.toString
+    args   := #["available-imports"]
+  }
+  let lakeProc ← IO.Process.spawn spawnArgs
+  let stdout := String.trim (← lakeProc.stdout.readToEnd)
+  let exitCode ← lakeProc.wait
+  match exitCode with
+  | 0 =>
+    let Except.ok (availableImports : AvailableImports) ← pure (Json.parse stdout >>= fromJson?)
+      | throw <| IO.userError  s!"invalid output from `lake available-imports`:\n{stdout}"
+    return availableImports
+  | _ =>
+    return none
+
+def collectAvailableImportsFromSrcSearchPath : IO AvailableImports :=
+  (·.2) <$> StateT.run (s := #[]) do
+    let srcSearchPath ← initSrcSearchPath
+    for p in srcSearchPath do
+      if ! (← p.isDir) then
+        continue
+      Lean.forEachModuleInDir p fun mod => do
+        modify (·.push mod)
+
+def collectAvailableImports : IO AvailableImports := do
+  match ← ImportCompletion.collectAvailableImportsFromLake with
+  | none =>
+    -- lake is not available => walk LEAN_SRC_PATH for an approximation
+    ImportCompletion.collectAvailableImportsFromSrcSearchPath
+  | some availableImports => pure availableImports
+
+def find (text : FileMap) (headerStx : Syntax) (params : CompletionParams) (availableImports : AvailableImports) : CompletionList :=
+  let availableImports := availableImports.toImportTrie
+  let completionPos := text.lspPosToUtf8Pos params.position
+  if isImportNameCompletionRequest headerStx completionPos then
+    let allAvailableImportNameCompletions := availableImports.toArray.map ({ label := toString · })
+    { isIncomplete := false, items := allAvailableImportNameCompletions }
+  else if isImportCmdCompletionRequest headerStx completionPos then
+    let allAvailableFullImportCompletions := availableImports.toArray.map ({ label := s!"import {·}" })
+    { isIncomplete := false, items := allAvailableFullImportCompletions }
+  else
+    let completionNames : Array Name := computePartialImportCompletions headerStx completionPos availableImports
+    let completions : Array CompletionItem := completionNames.map ({ label := toString · })
+    { isIncomplete := false, items := completions }
+
+def computeCompletions (text : FileMap) (headerStx : Syntax) (params : CompletionParams)
+    : IO CompletionList := do
+  let availableImports ← collectAvailableImports
+  return find text headerStx params availableImports
+
+end ImportCompletion
--- a/src/Lean/Server/Watchdog.lean
+++ b/src/Lean/Server/Watchdog.lean
@ -697,7 +697,7 @@ def loadReferences : IO References := do

 def initAndRunWatchdog (args : List String) (i o e : FS.Stream) : IO Unit := do
  let workerPath ← findWorkerPath
-  let srcSearchPath ← initSrcSearchPath (← getBuildDir)
+  let srcSearchPath ← initSrcSearchPath
  let references ← IO.mkRef (← loadReferences)
  let fileWorkersRef ← IO.mkRef (RBMap.empty : FileWorkerMap)
  let i ← maybeTee "wdIn.txt" false i
--- a/src/Lean/Util/LakePath.lean
+++ b/src/Lean/Util/LakePath.lean
@ -0,0 +1,15 @@
+/-
+Copyright (c) 2023 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Authors: Sebastian Ullrich
+-/
+
+def Lean.determineLakePath : IO System.FilePath := do
+  if let some lakePath ← IO.getEnv "LAKE" then
+    return System.FilePath.mk lakePath
+
+  let sysroot? ← IO.getEnv "LEAN_SYSROOT"
+  let lakePath ← match sysroot? with
+    | some sysroot => pure <| System.FilePath.mk sysroot / "bin" / "lake"
+    | none         => pure <| (← IO.appDir) / "lake"
--- a/src/Lean/Util/Path.lean
+++ b/src/Lean/Util/Path.lean
@ -13,6 +13,17 @@ import Lean.Data.Name
 namespace Lean
 open System

+partial def forEachModuleInDir [Monad m] [MonadLiftT IO m]
+    (dir : FilePath) (f : Lean.Name → m PUnit) (ext := "lean") : m PUnit := do
+  for entry in (← dir.readDir) do
+    if (← liftM (m := IO) <| entry.path.isDir) then
+      let n := Lean.Name.mkSimple entry.fileName
+      let r := FilePath.withExtension entry.fileName ext
+      if (← liftM (m := IO) r.pathExists) then f n
+      forEachModuleInDir entry.path (f <| n ++ ·)
+    else if entry.path.extension == some ext then
+      f <| Lean.Name.mkSimple <| FilePath.withExtension entry.fileName "" |>.toString
+
 def realPathNormalized (p : FilePath) : IO FilePath :=
  return (← IO.FS.realPath p).normalize

--- a/src/Lean/Util/Paths.lean
+++ b/src/Lean/Util/Paths.lean
@ -19,13 +19,12 @@ structure LeanPaths where
  loadDynlibPaths : Array FilePath := #[]
  deriving ToJson, FromJson

-def initSrcSearchPath (_leanSysroot : FilePath) (sp : SearchPath := ∅) : IO SearchPath := do
-  let srcSearchPath :=
-    if let some p := (← IO.getEnv "LEAN_SRC_PATH") then
-      System.SearchPath.parse p
-    else []
+def initSrcSearchPath (pkgSearchPath : SearchPath := ∅) : IO SearchPath := do
+  let srcSearchPath := (← IO.getEnv "LEAN_SRC_PATH")
+    |>.map System.SearchPath.parse
+    |>.getD []
  let srcPath := (← IO.appDir) / ".." / "src" / "lean"
  -- `lake/` should come first since on case-insensitive file systems, Lean thinks that `src/` also contains `Lake/`
-  return srcSearchPath ++ sp ++ [srcPath / "lake", srcPath]
+  return srcSearchPath ++ pkgSearchPath ++ [srcPath / "lake", srcPath]

 end Lean
--- a/src/lake/Lake/Config/Glob.lean
+++ b/src/lake/Lake/Config/Glob.lean
@ -23,17 +23,6 @@ deriving Inhabited, Repr

 instance : Coe Name Glob := ⟨Glob.one⟩

-partial def forEachModuleIn [Monad m] [MonadLiftT IO m]
-(dir : FilePath) (f : Name → m PUnit) (ext := "lean") : m PUnit := do
-  for entry in (← dir.readDir) do
-    if (← liftM (m := IO) <| entry.path.isDir) then
-      let n := Name.mkSimple entry.fileName
-      let r := FilePath.withExtension entry.fileName ext
-      if (← liftM (m := IO) r.pathExists) then f n
-      forEachModuleIn entry.path (f <| n ++ ·)
-    else if entry.path.extension == some ext then
-      f <| Name.mkSimple <| FilePath.withExtension entry.fileName "" |>.toString
-
 namespace Glob

 def «matches» (m : Name) : (self : Glob) → Bool
@ -45,6 +34,6 @@ def «matches» (m : Name) : (self : Glob) → Bool
 (dir : FilePath) (f : Name → m PUnit) : (self : Glob) → m PUnit
 | one n => f n
 | submodules n =>
-  forEachModuleIn (Lean.modToFilePath dir n "") (f <| n ++ ·)
+  Lean.forEachModuleInDir (Lean.modToFilePath dir n "") (f <| n ++ ·)
 | andSubmodules n =>
-  f n *> forEachModuleIn (Lean.modToFilePath dir n "") (f <| n ++ ·)
+  f n *> Lean.forEachModuleInDir (Lean.modToFilePath dir n "") (f <| n ++ ·)