From 72477f3cc6c81187ed7efda0d1ec8b33d46eff4a Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 11 Jul 2019 17:22:26 -0700
Subject: [PATCH] feat(library/init/lean/parser/command): improve how
 `optional` affects `firstTokens` field

---
 library/init/lean/parser/command.lean | 12 +-------
 library/init/lean/parser/parser.lean  | 41 ++++++++++++++++++---------
 tests/playground/cmdparsertest1.lean  |  1 +
 3 files changed, 30 insertions(+), 24 deletions(-)

diff --git a/library/init/lean/parser/command.lean b/library/init/lean/parser/command.lean
index 3ad7b9debf..b7b4a1374c 100644
--- a/library/init/lean/parser/command.lean
+++ b/library/init/lean/parser/command.lean
@@ -45,17 +45,7 @@ def declVal          := declValSimple <|> declValEqns
 def «def»            := parser! "def " >> declId >> optDeclSig >> declVal
 def «theorem»        := parser! "theorem " >> declId >> declSig >> declVal
 
-def declaration      := declModifiers >> («def» <|> «theorem»)
-
-/- TODO: remove this hack by improving how we compute `Parser.info.firstTokens` -/
-def declTokens := ["/--", "@[", "private", "protected", "noncomputable", "unsafe", "def", "theorem"]
-
-@[builtinCommandParser] def declEntry : Parser :=
-{ info := {
-    firstTokens := FirstTokens.tokens $ declTokens.map $ fun tk => { val := tk },
-    .. declaration.info
-  },
-  fn   := declaration.fn }
+@[builtinCommandParser] def declaration := declModifiers >> («def» <|> «theorem»)
 
 end Command
 
diff --git a/library/init/lean/parser/parser.lean b/library/init/lean/parser/parser.lean
index ec616394d0..d870581119 100644
--- a/library/init/lean/parser/parser.lean
+++ b/library/init/lean/parser/parser.lean
@@ -149,26 +149,37 @@ def ParserFn (k : ParserKind) := ParserArg k → BasicParserFn
 instance ParserFn.inhabited (k : ParserKind) : Inhabited (ParserFn k) := ⟨fun _ _ => id⟩
 
 inductive FirstTokens
-| epsilon : FirstTokens
-| unknown : FirstTokens
-| tokens  : List TokenConfig → FirstTokens
+| epsilon   : FirstTokens
+| unknown   : FirstTokens
+| tokens    : List TokenConfig → FirstTokens
+| optTokens : List TokenConfig → FirstTokens
 
 namespace FirstTokens
 
 def merge : FirstTokens → FirstTokens → FirstTokens
-| epsilon     tks         := tks
-| tks         epsilon     := tks
-| (tokens s₁) (tokens s₂) := tokens (s₁ ++ s₂)
-| _           _           := unknown
+| epsilon        tks            := tks
+| tks            epsilon        := tks
+| (tokens s₁)    (tokens s₂)    := tokens (s₁ ++ s₂)
+| (optTokens s₁) (optTokens s₂) := optTokens (s₁ ++ s₂)
+| (tokens s₁)    (optTokens s₂) := tokens (s₁ ++ s₂)
+| (optTokens s₁) (tokens s₂)    := tokens (s₁ ++ s₂)
+| _              _              := unknown
 
 def seq : FirstTokens → FirstTokens → FirstTokens
-| epsilon tks := tks
-| tks     _   := tks
+| epsilon        tks            := tks
+| (optTokens s₁) (optTokens s₂) := optTokens (s₁ ++ s₂)
+| (optTokens s₁) (tokens s₂)    := tokens (s₁ ++ s₂)
+| tks            _              := tks
+
+def toOptional : FirstTokens → FirstTokens
+| (tokens tks) := optTokens tks
+| tks          := tks
 
 def toStr : FirstTokens → String
-| epsilon      := "epsilon"
-| unknown      := "unknown"
-| (tokens tks) := toString tks
+| epsilon         := "epsilon"
+| unknown         := "unknown"
+| (tokens tks)    := toString tks
+| (optTokens tks) := "?" ++ toString tks
 
 instance : HasToString FirstTokens := ⟨toStr⟩
 
@@ -288,8 +299,12 @@ fun a c s =>
   let s      := if s.hasError && s.pos == iniPos then s.restore iniSz iniPos else s;
   s.mkNode nullKind iniSz
 
+@[noinline] def optionaInfo (p : ParserInfo) : ParserInfo :=
+{ updateTokens := p.updateTokens,
+  firstTokens  := p.firstTokens.toOptional }
+
 @[inline] def optional {k : ParserKind} (p : Parser k) : Parser k :=
-{ info := noFirstTokenInfo p.info,
+{ info := optionaInfo p.info,
   fn   := optionalFn p.fn }
 
 @[inline] def lookaheadFn {k : ParserKind} (p : ParserFn k) : ParserFn k :=
diff --git a/tests/playground/cmdparsertest1.lean b/tests/playground/cmdparsertest1.lean
index 22038acfa9..c31dda1799 100644
--- a/tests/playground/cmdparsertest1.lean
+++ b/tests/playground/cmdparsertest1.lean
@@ -16,6 +16,7 @@ is.mfor $ fun input => do
 
 def main (xs : List String) : IO Unit :=
 do
+IO.println Command.declaration.info.firstTokens;
 test [
 "@[inline] def x := 2",
 "protected def length.{u} {α : Type u} : List α → Nat