feat: use leantar in lake exe cache (leanprover-community#5710)

This hooks up the new `leantar` tool from https://github.com/digama0/leangz to `lake exe cache`. It uses an olean compressor to achieve much smaller file sizes, and faster decompression. See the [Zulip discussion](https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/leangz.3A.20olean.20file.20compressor/near/369738648).

The new files have a new file extension, `123.ltar` instead of `123.tar.gz`. This implementation does not multiplex between them, so it will be a hard reset. In fact, the new cache doesn't even know how to clear the old cache files, so you might need to do that manually. Switching between branches will still work during the interim, since the cache will have both kinds of file and use the right ones.

Author: digama0

Cache/IO.lean

@@ -7,6 +7,7 @@ Authors: Arthur Paulino
 import Lean.Data.HashMap
 import Lean.Data.RBMap
 import Lean.Data.RBTree
+import Lean.Data.Json.Printer
 
 /-- Removes a parent path from the beginning of a path -/
 def System.FilePath.withoutParent (path parent : FilePath) : FilePath :=
@@ -16,8 +17,15 @@ def System.FilePath.withoutParent (path parent : FilePath) : FilePath :=
     | x, [] => x
   mkFilePath $ aux path.components parent.components
 
-def UInt64.asTarGz (n : UInt64) : String :=
-  s!"{n}.tar.gz"
+def Nat.toHexDigits (n : Nat) : Nat → (res : String := "") → String
+  | 0, s => s
+  | len+1, s =>
+    let b := UInt8.ofNat (n >>> (len * 8))
+    Nat.toHexDigits n len <|
+      s.push (Nat.digitChar (b >>> 4).toNat) |>.push (Nat.digitChar (b &&& 15).toNat)
+
+def UInt64.asLTar (n : UInt64) : String :=
+  s!"{Nat.toHexDigits n.toNat 8}.ltar"
 
 namespace Cache.IO
 
@@ -51,12 +59,23 @@ def CURLBIN :=
   -- change file name if we ever need a more recent version to trigger re-download
   IO.CACHEDIR / s!"curl-{CURLVERSION}"
 
+/-- leantar version at https://github.com/digama0/leangz -/
+def LEANTARVERSION :=
+  "0.1.1"
+
+def LEANTARBIN :=
+  -- change file name if we ever need a more recent version to trigger re-download
+  IO.CACHEDIR / s!"leantar-{LEANTARVERSION}"
+
 def LAKEPACKAGESDIR : FilePath :=
   ⟨"lake-packages"⟩
 
 def getCurl : IO String := do
   return if (← CURLBIN.pathExists) then CURLBIN.toString else "curl"
 
+def getLeanTar : IO String := do
+  return if (← LEANTARBIN.pathExists) then LEANTARBIN.toString else "leantar"
+
 abbrev PackageDirs := Lean.RBMap String FilePath compare
 
 /-- Whether this is running on Mathlib repo or not -/
@@ -136,6 +155,44 @@ def validateCurl : IO Bool := do
     | _ => throw $ IO.userError "Invalidly formatted version of `curl`"
   | _ => throw $ IO.userError "Invalidly formatted response from `curl --version`"
 
+def Version := Nat × Nat × Nat
+  deriving Inhabited, DecidableEq
+
+instance : Ord Version := let _ := @lexOrd; lexOrd
+instance : LE Version := leOfOrd
+
+def parseVersion (s : String) : Option Version := do
+  let [maj, min, patch] := s.splitOn "." | none
+  some (maj.toNat!, min.toNat!, patch.toNat!)
+
+def validateLeanTar : IO Unit := do
+  if (← LEANTARBIN.pathExists) then return
+  if let some version ← some <$> runCmd "leantar" #["--version"] <|> pure none then
+    let "leantar" :: v :: _ := version.splitOn " "
+      | throw $ IO.userError "Invalidly formatted response from `leantar --version`"
+    let some v := parseVersion v | throw $ IO.userError "Invalidly formatted version of `leantar`"
+    -- currently we need exactly one version of leantar, change this to reflect compatibility
+    if v = (parseVersion LEANTARVERSION).get! then return
+  let win := System.Platform.getIsWindows ()
+  let target ← if win then
+    pure "x86_64-pc-windows-msvc"
+  else
+    let arch ← (·.trim) <$> runCmd "uname" #["-m"] false
+    unless arch ∈ ["x86_64", "aarch64"] do
+      throw $ IO.userError s!"unsupported architecture {arch}"
+    pure <|
+      if System.Platform.getIsOSX () then s!"{arch}-apple-darwin"
+      else s!"{arch}-unknown-linux-gnu"
+  IO.println s!"leantar is too old; downloading more recent version"
+  IO.FS.createDirAll IO.CACHEDIR
+  let ext := if win then "zip" else "tar.xz"
+  let _ ← runCmd "curl" #[
+    s!"https://github.com/digama0/leangz/releases/download/v{LEANTARVERSION}/leantar-v{LEANTARVERSION}-{target}.{ext}",
+    "-L", "-o", s!"{LEANTARBIN}.{ext}"]
+  let _ ← runCmd "tar" #["-xf", s!"{LEANTARBIN}.{ext}",
+    "-C", IO.CACHEDIR.toString, "--strip-components=1"]
+  let _ ← runCmd "mv" #[(IO.CACHEDIR / s!"leantar").toString, LEANTARBIN.toString]
+
 /-- Recursively gets all files from a directory with a certain extension -/
 partial def getFilesWithExtension
   (fp : FilePath) (extension : String) (acc : Array FilePath := #[]) :
@@ -150,7 +207,7 @@ namespace HashMap
 
 def filter (hashMap : HashMap) (set : Lean.RBTree String compare) (keep : Bool) : HashMap :=
   hashMap.fold (init := default) fun acc path hash =>
-    let contains := set.contains hash.asTarGz
+    let contains := set.contains hash.asLTar
     let add := if keep then contains else !contains
     if add then acc.insert path hash else acc
 
@@ -169,9 +226,9 @@ Each build file also has a `Bool` indicating whether that file is required for c
 def mkBuildPaths (path : FilePath) : IO $ Array (FilePath × Bool) := do
   let packageDir ← getPackageDir path
   return #[
+    (packageDir / LIBDIR / path.withExtension "trace", true),
     (packageDir / LIBDIR / path.withExtension "olean", true),
     (packageDir / LIBDIR / path.withExtension "ilean", true),
-    (packageDir / LIBDIR / path.withExtension "trace", true),
     (packageDir / IRDIR  / path.withExtension "c", true),
     (packageDir / LIBDIR / path.withExtension "extra", false)]
 
@@ -185,21 +242,25 @@ def allExist (paths : Array (FilePath × Bool)) : IO Bool := do
 def packCache (hashMap : HashMap) (overwrite : Bool) : IO $ Array String := do
   mkDir CACHEDIR
   IO.println "Compressing cache"
-  let mut acc := default
+  let mut acc := #[]
+  let mut tasks := #[]
   for (path, hash) in hashMap.toList do
-    let zip := hash.asTarGz
+    let zip := hash.asLTar
     let zipPath := CACHEDIR / zip
     let buildPaths ← mkBuildPaths path
     if ← allExist buildPaths then
       if overwrite || !(← zipPath.pathExists) then
-        discard $ runCmd "tar" $ #["-I", "gzip -9", "-cf", zipPath.toString] ++
-          ((← buildPaths.filterM (·.1.pathExists)) |>.map (·.1.toString))
+        tasks := tasks.push <| ← IO.asTask do
+          runCmd (← getLeanTar) $ #[zipPath.toString] ++
+            ((← buildPaths.filterM (·.1.pathExists)) |>.map (·.1.toString))
       acc := acc.push zip
+  for task in tasks do
+    _ ← IO.ofExcept task.get
   return acc
 
 /-- Gets the set of all cached files -/
 def getLocalCacheSet : IO $ Lean.RBTree String compare := do
-  let paths ← getFilesWithExtension CACHEDIR "gz"
+  let paths ← getFilesWithExtension CACHEDIR "ltar"
   return .fromList (paths.data.map (·.withoutParent CACHEDIR |>.toString)) _
 
 def isPathFromMathlib (path : FilePath) : Bool :=
@@ -215,19 +276,22 @@ def unpackCache (hashMap : HashMap) : IO Unit := do
   let hashMap := hashMap.filter (← getLocalCacheSet) true
   let size := hashMap.size
   if size > 0 then
+    let now ← IO.monoMsNow
     IO.println s!"Decompressing {size} file(s)"
     let isMathlibRoot ← isMathlibRoot
-    hashMap.forM fun path hash => do
-      let _ ← IO.asTask do
-      match path.parent with
-      | none | some path => do
-        let packageDir ← getPackageDir path
-        mkDir $ packageDir / LIBDIR / path
-        mkDir $ packageDir / IRDIR / path
+    let child ← IO.Process.spawn
+      { cmd := ← getLeanTar, args := #["-x", "-j", "-"], stdin := .piped }
+    let (stdin, child) ← child.takeStdin
+    let config : Array Lean.Json := hashMap.fold (init := #[]) fun config path hash =>
+      let pathStr := s!"{CACHEDIR / hash.asLTar}"
       if isMathlibRoot || !isPathFromMathlib path then
-        runCmd "tar" #["-xzf", s!"{CACHEDIR / hash.asTarGz}"]
+        config.push <| .str pathStr
       else -- only mathlib files, when not in the mathlib4 repo, need to be redirected
-        runCmd "tar" #["-xzf", s!"{CACHEDIR / hash.asTarGz}", "-C", mathlibDepPath.toString]
+        config.push <| .mkObj [("file", pathStr), ("base", mathlibDepPath.toString)]
+    stdin.putStr <| Lean.Json.compress <| .arr config
+    let exitCode ← child.wait
+    if exitCode != 0 then throw $ IO.userError s!"leantar failed with error code {exitCode}"
+    IO.println s!"unpacked in {(← IO.monoMsNow) - now} ms"
   else IO.println "No cache files to decompress"
 
 /-- Retrieves the azure token from the environment -/
@@ -241,7 +305,7 @@ instance : Ord FilePath where
 
 /-- Removes all cache files except for what's in the `keep` set -/
 def cleanCache (keep : Lean.RBTree FilePath compare := default) : IO Unit := do
-  for path in ← getFilesWithExtension CACHEDIR "gz" do
+  for path in ← getFilesWithExtension CACHEDIR "ltar" do
     if !keep.contains path then IO.FS.removeFile path
 
 end Cache.IO

Cache/Main.lean

@@ -53,11 +53,15 @@ def toPaths (args : List String) : List FilePath :=
 def curlArgs : List String :=
   ["get", "get!", "get-", "put", "put!", "commit", "commit!"]
 
+def leanTarArgs : List String :=
+  ["get", "get!", "pack", "pack!", "unpack"]
+
 open Cache IO Hashing Requests in
 def main (args : List String) : IO Unit := do
   let hashMemo ← getHashMemo
   let hashMap := hashMemo.hashMap
-  let goodCurl := !(curlArgs.contains (args.headD "") && !(← validateCurl))
+  let goodCurl ← pure !curlArgs.contains (args.headD "") <||> validateCurl
+  if leanTarArgs.contains (args.headD "") then validateLeanTar
   match args with
   | ["get"] => getFiles hashMap false goodCurl true
   | ["get!"] => getFiles hashMap true goodCurl true
@@ -69,7 +73,7 @@ def main (args : List String) : IO Unit := do
   | ["pack!"] => discard $ packCache hashMap true
   | ["unpack"] => unpackCache hashMap
   | ["clean"] =>
-    cleanCache $ hashMap.fold (fun acc _ hash => acc.insert $ CACHEDIR / hash.asTarGz) .empty
+    cleanCache $ hashMap.fold (fun acc _ hash => acc.insert $ CACHEDIR / hash.asLTar) .empty
   | ["clean!"] => cleanCache
   | ["put"] => putFiles (← packCache hashMap false) false (← getToken)
   | ["put!"] => putFiles (← packCache hashMap false) true (← getToken)

Cache/Requests.lean

@@ -31,7 +31,7 @@ def mkGetConfigContent (hashMap : IO.HashMap) : IO String := do
   -- We sort the list so that the large files in `MathlibExtras` are requested first.
   hashMap.toArray.qsort (fun ⟨p₁, _⟩ ⟨_, _⟩ => p₁.components.head? = "MathlibExtras")
     |>.foldlM (init := "") fun acc ⟨_, hash⟩ => do
-    let fileName := hash.asTarGz
+    let fileName := hash.asLTar
     -- Below we use `String.quote`, which is intended for quoting for use in Lean code
     -- this does not exactly match the requirements for quoting for curl:
     -- ```
@@ -47,7 +47,7 @@ def mkGetConfigContent (hashMap : IO.HashMap) : IO String := do
 
 /-- Calls `curl` to download a single file from the server to `CACHEDIR` (`.cache`) -/
 def downloadFile (hash : UInt64) : IO Bool := do
-  let fileName := hash.asTarGz
+  let fileName := hash.asLTar
   let url ← mkFileURL fileName none
   let path := IO.CACHEDIR / fileName
   let out ← IO.Process.output