From 6a67c13044e89ba8ae5ebfa9915b0af69d332392 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 11 Aug 2022 17:37:50 -0700
Subject: [PATCH] feat: generalize helper functions

---
 src/Lean/Compiler/CompilerM.lean | 25 +++--------------------
 src/Lean/Compiler/InferType.lean | 35 ++++++++++++++++++++++++++++++--
 src/Lean/Compiler/LCNF.lean      |  5 ++---
 3 files changed, 38 insertions(+), 27 deletions(-)

diff --git a/src/Lean/Compiler/CompilerM.lean b/src/Lean/Compiler/CompilerM.lean
index 830d544138..9930011c9f 100644
--- a/src/Lean/Compiler/CompilerM.lean
+++ b/src/Lean/Compiler/CompilerM.lean
@@ -23,28 +23,9 @@ instance : AddMessageContext CompilerM where
     let opts ← getOptions
     return MessageData.withContext { env, lctx, opts, mctx := {} } msgData
 
-/--
-Infer the type of a LCNF expression.
--/
-def inferType (e : Expr) : CompilerM Expr := do
-  InferType.inferType e { lctx := (← get).lctx }
-
-def getLevel (type : Expr) : CompilerM Level := do
-  match (← inferType type) with
-  | .sort u => return u
-  | e => if e.isAnyType then return levelOne else throwError "type expected{indentExpr type}"
-
-/-- Create `lcUnreachable type` -/
-def mkLcUnreachable (type : Expr) : CompilerM Expr := do
-  let u ← getLevel type
-  return .app (.const ``lcUnreachable [u]) type
-
-/-- Create `lcCast expectedType e : expectedType` -/
-def mkLcCast (e : Expr) (expectedType : Expr) : CompilerM Expr := do
-  let type ← inferType e
-  let u ← getLevel type
-  let v ← getLevel expectedType
-  return mkApp3 (.const ``lcCast [u, v]) type expectedType e
+instance : MonadInferType CompilerM where
+  inferType e := do
+    InferType.inferType e { lctx := (← get).lctx }
 
 def mkLocalDecl (binderName : Name) (type : Expr) (bi := BinderInfo.default) : CompilerM Expr := do
   let fvarId ← mkFreshFVarId
diff --git a/src/Lean/Compiler/InferType.lean b/src/Lean/Compiler/InferType.lean
index 8b7e75beea..38f824eda5 100644
--- a/src/Lean/Compiler/InferType.lean
+++ b/src/Lean/Compiler/InferType.lean
@@ -6,7 +6,12 @@ Authors: Leonardo de Moura
 import Lean.Compiler.LCNFTypes
 import Lean.Compiler.Util
 
-namespace Lean.Compiler.InferType
+namespace Lean.Compiler
+
+class MonadInferType (m : Type → Type) where
+  inferType : Expr → m Expr
+
+namespace InferType
 
 structure InferTypeM.Context where
   lctx : LocalContext
@@ -140,4 +145,30 @@ mutual
     | .forallE ..    => inferForallType e
     | .lam ..        => inferLambdaType e
     | .letE ..       => inferLambdaType e
-end
\ No newline at end of file
+end
+
+end InferType
+
+instance : MonadInferType InferType.InferTypeM where
+  inferType := InferType.inferType
+
+export MonadInferType (inferType)
+
+def getLevel [Monad m] [MonadInferType m] [MonadError m] (type : Expr) : m Level := do
+  match (← inferType type) with
+  | .sort u => return u
+  | e => if e.isAnyType then return levelOne else throwError "type expected{indentExpr type}"
+
+/-- Create `lcUnreachable type` -/
+def mkLcUnreachable [Monad m] [MonadInferType m] [MonadError m] (type : Expr) : m Expr := do
+  let u ← getLevel type
+  return .app (.const ``lcUnreachable [u]) type
+
+/-- Create `lcCast expectedType e : expectedType` -/
+def mkLcCast [Monad m] [MonadInferType m] [MonadError m] (e : Expr) (expectedType : Expr) : m Expr := do
+  let type ← inferType e
+  let u ← getLevel type
+  let v ← getLevel expectedType
+  return mkApp3 (.const ``lcCast [u, v]) type expectedType e
+
+end Lean.Compiler
\ No newline at end of file
diff --git a/src/Lean/Compiler/LCNF.lean b/src/Lean/Compiler/LCNF.lean
index 101cc5dbb0..61168d6189 100644
--- a/src/Lean/Compiler/LCNF.lean
+++ b/src/Lean/Compiler/LCNF.lean
@@ -41,9 +41,8 @@ structure State where
 
 abbrev M := ReaderT Context $ StateRefT State CoreM
 
-/-- Return the type of the given LCNF expression. -/
-def inferType (e : Expr) : M Expr := do
-  InferType.inferType e { lctx := (← get).lctx' }
+instance : MonadInferType M where
+  inferType e := do InferType.inferType e { lctx := (← get).lctx' }
 
 @[inline] def liftMetaM (x : MetaM α) : M α := do
   x.run' { lctx := (← get).lctx }