File size: 3,436 Bytes
dfe1826 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | import Mathlib
/-!
Dump the public Mathlib theorems of the elaborated environment as JSON Lines.
Run from `lean_project/` (done by `scripts/build_index.py`):
FORMAL_MAS_CORPUS_OUT=<file> [FORMAL_MAS_MODULE_PREFIX=Mathlib.Foo] lake env lean <this file>
Line 1 is `{"_meta": ...}`; every other line is one theorem with its raw signature,
docstring, declaration range and `to_additive`/`to_dual` target flags. Hard filters that
need the environment (private, internal, deprecated, instances) are applied here; origin
classification and the remaining filters live in `formal_mas.retrieval.corpus`.
-/
open Lean Meta Elab Command
set_option pp.universes false
set_option format.width 100000
#eval show CommandElabM Unit from do
let some out ← IO.getEnv "FORMAL_MAS_CORPUS_OUT"
| throwError "set FORMAL_MAS_CORPUS_OUT to the output path"
let pfx := ((← IO.getEnv "FORMAL_MAS_MODULE_PREFIX").getD "Mathlib").toName
let env ← getEnv
let opts ← getOptions
let inScope (n : Name) : Option Name := do
let idx ← env.getModuleIdxFor? n
let mod := env.header.moduleNames[idx.toNat]!
if pfx.isPrefixOf mod then some mod else none
-- Pass 1: names produced by the translation attributes (targets of the maps).
let mut additive : NameSet := {}
let mut dual : NameSet := {}
for (n, ci) in env.constants.map₁.toList do
let .thmInfo _ := ci | continue
if let some t := Mathlib.Tactic.ToAdditive.translations.find? env n then
additive := additive.insert t.translation
if let some t := Mathlib.Tactic.ToDual.translations.find? env n then
dual := dual.insert t.translation
-- Pass 2: emit candidates.
let h ← IO.FS.Handle.mk out .write
h.putStrLn <| Json.compress <| Json.mkObj [("_meta", Json.mkObj [
("lean_version", toJson Lean.versionString),
("module_prefix", toJson pfx.toString),
("pp_function", toJson "Lean.PrettyPrinter.ppSignature"),
("pp_options", Json.mkObj [
("pp.universes", toJson (pp.universes.get opts)),
("format.width", toJson (Std.Format.format.width.get opts))])])]
let mut emitted : Nat := 0
let mut skipped : Nat := 0
let mut failed : Nat := 0
for (n, ci) in env.constants.map₁.toList do
let .thmInfo _ := ci | continue
let some mod := inScope n | continue
if isPrivateName n || n.isInternalDetail || Linter.isDeprecated env n
|| isInstanceCore env n then
skipped := skipped + 1
continue
let sig? ← liftTermElabM do
try return some (toString (← PrettyPrinter.ppSignature n).fmt)
catch _ => return none
let some sig := sig? | failed := failed + 1; continue
let doc ← findDocString? env n
let range : Json := match ← findDeclarationRanges? n with
| some r => toJson [r.selectionRange.pos.line, r.selectionRange.pos.column,
r.selectionRange.endPos.line, r.selectionRange.endPos.column]
| none => Json.null
h.putStrLn <| Json.compress <| Json.mkObj [
("name", toJson n.toString), ("module", toJson mod.toString),
("signature_raw", toJson sig), ("docstring", toJson doc), ("range", range),
("to_additive", toJson (additive.contains n)), ("to_dual", toJson (dual.contains n))]
emitted := emitted + 1
if emitted % 20000 == 0 then IO.eprintln s!"extract_corpus: {emitted} emitted"
h.flush
logInfo m!"extract_corpus: emitted {emitted}, skipped {skipped}, pp failures {failed}"
|