From bf71068b145e77a4e4f9cbeac750fdb041b02d91 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 2 May 2018 16:59:50 -0700
Subject: [PATCH] feat(library/init/lean/ir): parse IR definitions

---
 library/init/lean/ir/check.lean    | 18 ++++++++----
 library/init/lean/ir/format.lean   | 13 ++++++---
 library/init/lean/ir/ir.lean       |  8 ++++--
 library/init/lean/ir/parser.lean   | 46 ++++++++++++++++++++++++++++++
 library/init/lean/ir/reserved.lean |  6 +++-
 tests/lean/run/parser_ir1.lean     | 18 ++++++++++++
 6 files changed, 96 insertions(+), 13 deletions(-)
 create mode 100644 tests/lean/run/parser_ir1.lean

diff --git a/library/init/lean/ir/check.lean b/library/init/lean/ir/check.lean
index 1e4f19fd61..caeb9c1c3d 100644
--- a/library/init/lean/ir/check.lean
+++ b/library/init/lean/ir/check.lean
@@ -61,13 +61,14 @@ def arg.declare_at (b : blockid) : arg → ssa_decl_m unit
 /- Collect where each variable is declared, and
    check whether each variable was declared at most once. -/
 def decl.declare_vars : decl → ssa_decl_m unit
-| {as := as, bs := b::bs, ..} :=
+| (decl.defn {args := as, ..} (b::bs)) :=
   /- We assume that arguments are declared in the first basic block.
      TODO: check whether this assumption matches LLVM or not -/
   as.mmap' (arg.declare_at b.id) >>
   b.declare_vars >>
   bs.mmap' block.declare_vars
-| _ := throw ssa_error.no_block
+| (decl.defn _  []) := throw ssa_error.no_block
+| _                 := return ()
 
 /- Generate the mapping from variable to blockid for the given declaration.
    This function assumes `d` is in SSA. -/
@@ -173,8 +174,12 @@ defined before being used.
 -/
 def decl.valid_ssa (d : decl) : except_t ssa_error id var2blockid :=
 do m ← d.var2blockid,
-   d.bs.mmap' (λ b : block, run_reader block.valid_ssa (m, b)),
-   return m
+   match d with
+   | decl.defn _ bs := do
+      bs.mmap' (λ b : block, run_reader block.valid_ssa (m, b)),
+      return m
+   | _ := return m
+   end
 
 /- Check blockids -/
 inductive blockid_error
@@ -202,8 +207,9 @@ def terminator.check_blockids : terminator → blockid_check_m unit
 def block.check_blockids (b : block) : blockid_check_m unit :=
 b.term.check_blockids
 
-def decl.check_blockids (d : decl) : blockid_check_m unit :=
-d.bs.mmap' block.declare >> d.bs.mmap' block.check_blockids
+def decl.check_blockids : decl → blockid_check_m unit
+| (decl.defn _ bs) := bs.mmap' block.declare >> bs.mmap' block.check_blockids
+| _                := return ()
 
 def check_blockids (d : decl) : except_t blockid_error id blockid_set :=
 run_state (d.check_blockids >> get) mk_blockid_set
diff --git a/library/init/lean/ir/format.lean b/library/init/lean/ir/format.lean
index bae501edcd..92c66c822b 100644
--- a/library/init/lean/ir/format.lean
+++ b/library/init/lean/ir/format.lean
@@ -117,7 +117,7 @@ instance terminator.has_to_string : has_to_string terminator := ⟨pretty ∘ to
 
 def block.to_format : block → format
 | {id := id, phis := ps, instrs := is, term := t} :=
-  to_fmt id ++ ":" ++ line ++ join_sep ps line ++ join_sep is line ++ to_fmt t
+  to_fmt id ++ ":" ++ nest 2 (line ++ join_suffix ps (";" ++ line) ++ join_suffix is (";" ++ line) ++ to_fmt t ++ ";")
 
 instance block.has_to_format : has_to_format block := ⟨block.to_format⟩
 instance block.has_to_string : has_to_string block := ⟨pretty ∘ to_fmt⟩
@@ -134,10 +134,15 @@ def result.to_format : result → format
 instance result.has_to_format : has_to_format result := ⟨result.to_format⟩
 instance result.has_to_string : has_to_string result := ⟨pretty ∘ to_fmt⟩
 
+def header.to_format (h : header) : format :=
+to_fmt h.n ++ " " ++ join_sep h.args " " ++ " : " ++ join_sep h.return " "
+
+instance header.has_to_format : has_to_format header := ⟨header.to_format⟩
+instance header.has_to_string : has_to_string header := ⟨pretty ∘ to_fmt⟩
+
 def decl.to_format : decl → format
-| {n := f, as := as, rs := rs, bs := bs} :=
-  "decl " ++ to_fmt f ++ " " ++ join_sep as " " ++ " : " ++ join_sep rs " " ++ " := " ++ line ++
-  join_sep bs line ++ line ++ "end"
+| (decl.defn h bs)  := "def " ++ to_fmt h ++ " := " ++ line ++ join_sep bs line ++ line
+| (decl.external h) := "external " ++ to_fmt h
 
 instance decl.has_to_format : has_to_format decl := ⟨decl.to_format⟩
 instance decl.has_to_string : has_to_string decl := ⟨pretty ∘ to_fmt⟩
diff --git a/library/init/lean/ir/ir.lean b/library/init/lean/ir/ir.lean
index 0d736bc6bc..980fa38237 100644
--- a/library/init/lean/ir/ir.lean
+++ b/library/init/lean/ir/ir.lean
@@ -102,8 +102,12 @@ structure arg :=
 structure result :=
 (ty : type)
 
-structure decl :=
-(n : fnid) (as : list arg) (rs : list result) (bs : list block)
+structure header :=
+(n : fnid) (args : list arg) (return : list result)
+
+inductive decl
+| external (h : header)
+| defn     (h : header) (bs : list block)
 
 end ir
 end lean
diff --git a/library/init/lean/ir/parser.lean b/library/init/lean/ir/parser.lean
index f0345351ca..33b374b180 100644
--- a/library/init/lean/ir/parser.lean
+++ b/library/init/lean/ir/parser.lean
@@ -6,6 +6,7 @@ Authors: Leonardo de Moura
 prelude
 import init.lean.ir.ir init.lean.parser.parser
 import init.lean.parser.identifier init.lean.parser.string_literal
+import init.lean.ir.reserved
 
 namespace lean
 namespace ir
@@ -73,12 +74,22 @@ try (do p ← pos,
         return $ uint16.of_nat n)
 <?> "uint16"
 
+def identifier : parser name :=
+try (do p ← pos,
+        n ← lean.parser.identifier,
+        when (is_reserved_name n) $ unexpected_at "keyword" p,
+        return n)
+<?> "identifier"
+
 def parse_var : parser var :=
 lexeme identifier <?> "variable"
 
 def parse_fnid : parser fnid :=
 lexeme identifier <?> "function name"
 
+def parse_blockid : parser blockid :=
+lexeme identifier <?> "label"
+
 def parse_nary_call (x : var) : parser instr :=
 do xs  ← many1 parse_var,
    symbol ":=",
@@ -127,5 +138,40 @@ def parse_instr : parser instr :=
 <|> (keyword "free" >> instr.free <$> parse_var)
 <|> parse_assignment
 
+def parse_phi : parser phi :=
+do (x, ty) ← try $ do { x ← parse_var, symbol ":", ty ← parse_type, symbol ":=", keyword "phi", return (x, ty) },
+   ys ← many1 parse_var,
+   return {x := x, ty := ty, ys := ys}
+
+def parse_terminator : parser terminator :=
+    (keyword "jmp" >> terminator.jmp <$> parse_blockid)
+<|> (keyword "ret" >> terminator.ret <$> many parse_var)
+<|> (keyword "case" >> terminator.case <$> parse_var <*> many parse_blockid)
+
+def parse_block : parser block :=
+do id ← parse_blockid,
+   symbol ":",
+   ps ← many (parse_phi <* symbol ";"),
+   is ← many (parse_instr <* symbol ";"),
+   t  ← parse_terminator <* symbol ";",
+   return { id := id, phis := ps, instrs := is, term := t }
+
+def parse_arg : parser arg :=
+do symbol "(", x ← parse_var, symbol ":", ty ← parse_type, symbol ")", return {n := x, ty := ty}
+
+
+def parse_header : parser header :=
+do n ← parse_fnid,
+   as ← many parse_arg,
+   symbol ":",
+   r ← many (result.mk <$> parse_type),
+   return { n := n, args := as, return := r }
+
+def parse_def : parser decl :=
+symbol "def" >> decl.defn <$> parse_header <*> (symbol ":=" >> many parse_block)
+
+def parse_external : parser decl :=
+symbol "external" >> decl.external <$> parse_header
+
 end ir
 end lean
diff --git a/library/init/lean/ir/reserved.lean b/library/init/lean/ir/reserved.lean
index 14e0646c00..0751730b59 100644
--- a/library/init/lean/ir/reserved.lean
+++ b/library/init/lean/ir/reserved.lean
@@ -16,7 +16,7 @@ def reserved := [ "bool", "byte", "uint16", "uint32", "uint64",
  "or", "xor", "le", "ge", "lt", "gt", "eq", "ne", "call", "closure",
  "apply", "cnstr", "set", "get", "sset", "sget", "array", "read",
  "write", "sarray", "sread", "swrite", "inc", "decs", "dec", "del",
- "phi", "ret", "case", "jmp", "decl", "end", "tt", "ff"]
+ "phi", "ret", "case", "jmp", "decl", "tt", "ff"]
 
 def reserved_set : rbtree string (<) :=
 reserved.foldl rbtree.insert (mk_rbtree string (<))
@@ -24,5 +24,9 @@ reserved.foldl rbtree.insert (mk_rbtree string (<))
 def is_reserved (s : string) : bool :=
 reserved_set.contains s
 
+def is_reserved_name : name → bool
+| (name.mk_string p s) := is_reserved s
+| _                    := ff
+
 end ir
 end lean
diff --git a/tests/lean/run/parser_ir1.lean b/tests/lean/run/parser_ir1.lean
new file mode 100644
index 0000000000..7ae61bc66b
--- /dev/null
+++ b/tests/lean/run/parser_ir1.lean
@@ -0,0 +1,18 @@
+import system.io
+import init.lean.ir.parser init.lean.ir.format
+
+open lean.parser
+open lean.ir
+
+def show_result {α} [lean.has_to_format α] (p : parser α) (s : string) : io unit :=
+match parse p s with
+| except.ok a    := io.print_ln (lean.to_fmt a)
+| except.error e := io.print_ln (e.to_string s)
+end
+
+def IR := "
+def succ (x : uint32) : uint32 :=
+main: one : uint32 := 1; x1 : uint32 := add x one; ret x1;
+"
+
+#eval show_result (whitespace >> parse_def) IR