| 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}" |
|
|
