RPINF: more detailed stats

Aesop/RPINF.lean

@@ -116,29 +116,39 @@ where
     else
       return acc
 
+structure PINFStats where
+  getArgTypes : Nanos := 0
+  inferType : Nanos := 0
+  isProp : Nanos := 0
+  whnf : Nanos := 0
+  instantiateForall : Nanos := 0
+  telescope : Nanos := 0
+  mkFVars : Nanos := 0
+
+instance : ToString PINFStats where
+  toString s := s!"getArgTypes: {s.getArgTypes.printAsMillis}, inferType: {s.inferType.printAsMillis}, isProp: {s.isProp.printAsMillis}, whnf: {s.whnf.printAsMillis}, instantiateForall: {s.instantiateForall.printAsMillis}, mkFVars: {s.mkFVars.printAsMillis}"
+
 variable [Monad m] [MonadRPINF m] [MonadLiftT MetaM m] [MonadControlT MetaM m]
   [MonadMCtx m] [MonadLiftT (ST IO.RealWorld) m] [MonadError m] [MonadRecDepth m]
   [MonadLiftT BaseIO m]
 
 -- `type?` may be `none` iff `e` is a type.
 @[specialize]
-partial def pinfCore (statsRef : IO.Ref Nanos) (e : Expr) (type? : Option Expr) : m Expr :=
+partial def pinfCore (statsRef : IO.Ref PINFStats) (e : Expr) (type? : Option Expr) : m Expr :=
   withIncRecDepth do
   checkCache e λ _ => do
-    if e.hasLooseBVars then
-      throwError "loose bvar in e:{indentD $ toString e}"
-    if let some type := type? then
-      if type.hasLooseBVars then
-        throwError "loose bvar in type:{indentD $ toString type}"
-    let e ← whnf e
+    let (e, ti) ← time $ whnf e
+    statsRef.modify λ s => { s with whnf := s.whnf + ti }
     match e with
     | .app .. =>
       ifNotProof do
         let f := e.getAppFn'
-        let fType ← inferType f
+        let (fType, ti) ← time $ inferType f
+        statsRef.modify λ s => { s with inferType := s.inferType + ti }
         let f ← pinfCore statsRef f fType
         let mut args := e.getAppArgs'
-        let argTypes ← getArgTypes fType args
+        let (argTypes, ti) ← time $ getArgTypes fType args
+        statsRef.modify λ s => { s with getArgTypes := s.getArgTypes + ti }
         for i in [:args.size] do
           let arg := args[i]!
           args := args.set! i default -- prevent nonlinear access to args[i]
@@ -150,18 +160,38 @@ partial def pinfCore (statsRef : IO.Ref Nanos) (e : Expr) (type? : Option Expr)
           return mkAppN f args
     | .lam .. =>
       -- TODO disable cache?
-      lambdaTelescope e λ xs e => withNewFVars xs do
-        let eType? ← type?.mapM λ type =>
-          withAtLeastTransparency .default $ instantiateForall type xs
-        mkLambdaFVars xs (← pinfCore statsRef e eType?)
+      let tiStart ← IO.monoNanosNow
+      lambdaTelescope e λ xs e => do
+        let tiEnd ← IO.monoNanosNow
+        statsRef.modify λ s => { s with telescope := s.telescope + ⟨tiEnd - tiStart⟩ }
+        withNewFVars xs do
+        let eType? ← type?.mapM λ type => do
+          let (type, ti) ← time do
+            withAtLeastTransparency .default $ instantiateForall type xs
+          statsRef.modify λ s =>
+            { s with instantiateForall := s.instantiateForall + ti }
+          return type
+        let e ← pinfCore statsRef e eType?
+        let (result, ti) ← time $ mkLambdaFVars xs e
+        statsRef.modify λ s => { s with mkFVars := s.mkFVars + ti }
+        return result
     | .forallE .. =>
       -- TODO disable cache?
       -- A `forall` is necessarily a type, and so is its conclusion.
-      forallTelescope e λ xs e => withNewFVars xs do
-        mkForallFVars xs (← pinfCore statsRef e none)
+      let tiStart ← IO.monoNanosNow
+      forallTelescope e λ xs e => do
+        let tiEnd ← IO.monoNanosNow
+        statsRef.modify λ s => { s with telescope := s.telescope + ⟨tiEnd - tiStart⟩ }
+        withNewFVars xs do
+        let e ← pinfCore statsRef e none
+        let (result, ti) ← time $ mkForallFVars xs e
+        statsRef.modify λ s => { s with mkFVars := s.mkFVars + ti }
+        return result
     | .proj t i e =>
       ifNotProof do
-        return .proj t i (← pinfCore statsRef e (← inferType e))
+        let (type, ti) ← time $ inferType e
+        statsRef.modify λ s => { s with inferType := s.inferType + ti }
+        return .proj t i (← pinfCore statsRef e type)
     | .mvar .. | .const .. | .fvar .. =>
       ifNotProof do
         return e
@@ -179,17 +209,21 @@ where
       lctx := lctx.modifyLocalDecl fvarId λ _ => ldecl
     withLCtx lctx (← getLocalInstances) k
 
-  isPropD (e : Expr) : MetaM Bool :=
+  @[inline, always_inline]
+  isPropD (e : Expr) : MetaM Bool := do
     -- withAtLeastTransparency .default $ isProp e -- TODO test perf impact
-    isProp e
+    let (result, ti) ← time $ isProp e
+    statsRef.modify λ s => { s with isProp := s.isProp + ti }
+    return result
 
+  @[inline, always_inline]
   ifNotProof (k : m Expr) : m Expr := do
     if let some type := type? then
       if ← isPropD type then
         return .mdata (mdataSetIsProof {}) e
     k
 
-def pinf (statsRef : IO.Ref Nanos) (e : Expr) : m Expr := do
+def pinf (statsRef : IO.Ref PINFStats) (e : Expr) : m Expr := do
   pinfCore statsRef (← instantiateMVars e) (← getType? e)
 where
   -- Returns the type of `e`, or `none` if `e` is a type.
@@ -203,20 +237,20 @@ where
       | .sort .. => pure none
       | _ => pure $ some type
 
-def pinf' (statsRef : IO.Ref Nanos) (e : Expr) : MetaM Expr := do
+def pinf' (statsRef : IO.Ref PINFStats) (e : Expr) : MetaM Expr := do
   (pinf statsRef e : RPINFT MetaM _).run' {}
 
-def rpinfExpr (statsRef : IO.Ref Nanos) (e : Expr) : m Expr :=
+def rpinfExpr (statsRef : IO.Ref PINFStats) (e : Expr) : m Expr :=
   withReducible $ pinf statsRef e
 
-def rpinfExpr' (statsRef : IO.Ref Nanos) (e : Expr) : MetaM Expr :=
+def rpinfExpr' (statsRef : IO.Ref PINFStats) (e : Expr) : MetaM Expr :=
   (rpinfExpr statsRef e : RPINFT MetaM _).run' {}
 
-def rpinf (statsRef : IO.Ref Nanos) (e : Expr) : m RPINF := do
+def rpinf (statsRef : IO.Ref PINFStats) (e : Expr) : m RPINF := do
   let expr ← rpinfExpr statsRef e
   return { expr, hash := rpinfHash expr }
 
-def rpinf' (statsRef : IO.Ref Nanos) (e : Expr) : MetaM RPINF :=
+def rpinf' (statsRef : IO.Ref PINFStats) (e : Expr) : MetaM RPINF :=
   (rpinf statsRef e : RPINFT MetaM _).run' {}
 
 end Aesop