Cryptic "unknown constant" error when using syntax extension

Stream: metaprogramming / tactics

[Dylan Ede]

Hello! I'm new to Lean but I've dived straight into metaprogramming. I'm interested in implementing "staged compilation" in the style of Toatie (github). To this end I've tried to make some syntax extensions to emulate the staging part of Toatie's type system. All was going well, until I hit a cryptic error message as detailed in the code below, which appears to be caused by the use of my new syntax, since a desugared version of the offending function (also presented below) does not encounter the error. Does anyone have any idea what is going on?

Note that the code is necessarily split into two modules since environment extensions cannot be used in the same module they are defined in.

LeanHDL/StageExt.lean

import Lean

open Lean EnvExtension

-- This extra environment state tracks the current syntactic staging "depth"
initialize stageEnv : EnvExtension Nat ← registerEnvExtension $ pure 0

def getStage [Monad m] [MonadEnv m] : m Nat := stageEnv.getState <$> getEnv

-- Executes `x` in an environment with the specified `stageEnv` value
def withStage [Monad m] [MonadFinally m] [MonadEnv m] (s : Nat) (x : m α) : m α := do
    withEnv (stageEnv.setState (← getEnv) s) x

LeanHDL/Basic.lean

import Lean
import LeanHDL.StageExt

/-
  (Very) partial re-implementation in Lean of "Toatie: functional hardware description with dependent types" by Craig Ramsay
  https://stax.strath.ac.uk/concern/theses/6w924c39c
-/

-- Unlike in the paper, I've put stage tracking in the `Code` type itself
-- Ideally there should be no way of using this type except through the syntax defined below
structure Code (stage : Nat) (α : Sort u) where
  «~» : α

section Syntax
  open Lean Elab Term

  -- The corresponding syntax in the toatie paper is ⟨_⟩
  syntax (name := «stx ⟪_⟫») "⟪" term "⟫" : term
  @[term_elab «stx ⟪_⟫»]
  def «elab ⟪_⟫» : TermElab
    | `(⟪ $ty:term ⟫), ty? => do
      let stage <- getStage
      let stageLit := Lean.Syntax.mkNumLit $ toString stage
      withStage stage.succ $ elabTerm (<- `(Code $stageLit:num $ty)) ty?
    | _, _ => throwUnsupportedSyntax

  -- The corresponding syntax in the toatie paper is ⟦_⟧
  syntax (name := «stx [|_|]»)"[|" term "|]" : term
  @[term_elab «stx [|_|]»]
  def «elab [|_|]» : TermElab
    | `([| $t |]), ty? => do
      let stage <- getStage
      let stageLit := Lean.Syntax.mkNumLit $ toString stage
      withStage stage.succ $ elabTerm (<- `(Code.mk (stage := $stageLit:num) $t)) ty?
    | _, _ => throwUnsupportedSyntax

  -- The corresponding syntax in the toatie paper is ˜_
  syntax:max (name := «stx ~_») "~" term:max : term
  @[term_elab «stx ~_»]
  def «elab ~_» : TermElab
    | `(~ $t), ty? => do
      let stage <- getStage
      let stagePredLit := Lean.Syntax.mkNumLit $ toString stage.pred
      let «⟪t⟫» <- withStage stage.pred $ elabTerm (<- `(Code.«~» (stage := $stagePredLit:num) $t)) ty?
      if stage == Nat.zero && «⟪t⟫».hasFVar then
        throwError "Free variables not allowed when escaping to stage 0"
      pure «⟪t⟫»
    | _, _ => throwUnsupportedSyntax

-- All of the following types and functions are direct translations from the examples in the paper

inductive Bit : Nat -> Type where
  | O : Bit 0
  | I : Bit 1

inductive Unsigned : Nat -> Nat -> Type where
  | nil : Unsigned 0 0
  | cons : Unsigned width val -> Bit b -> Unsigned width.succ (b + val * 2)

inductive BitPair : Nat -> Type where
  | mk : Bit a -> Bit b -> {p : c = a * 2 + b} -> BitPair c

def bitAdd : Bit c -> Bit x -> Bit y -> BitPair (x + y + c)
  | .O, .O, .O => BitPair.mk .O .O (p := rfl)
  | .O, .O, .I => BitPair.mk .O .I (p := rfl)
  | .O, .I, .O => BitPair.mk .O .I (p := rfl)
  | .O, .I, .I => BitPair.mk .I .O (p := rfl)
  | .I, .O, .O => BitPair.mk .O .I (p := rfl)
  | .I, .O, .I => BitPair.mk .I .O (p := rfl)
  | .I, .I, .O => BitPair.mk .I .O (p := rfl)
  | .I, .I, .I => BitPair.mk .I .I (p := rfl)

/- Why does `addU` fail to compile with (only) the below error message?

   unknown constant 'addU.match_2'
-/
def addU : {w x y : Nat} -> ⟪Unsigned w x⟫ -> ⟪Unsigned w y⟫ -> ⟪Bit c⟫ -> ⟪Unsigned w.succ (x + y + c)⟫
  | 0, 0, 0, _, _, cin => cast (by simp) [|Unsigned.cons .nil ~cin|]
  | w + 1, _, _, ux, uy, cin => [|
    let .cons xbs xb := ~ux
    let .cons ybs yb := ~uy
    let .mk cin' lsb := bitAdd ~cin xb yb
    let r := ~(addU [|xbs|] [|ybs|] [|cin'|])
    cast (by congr 1; simp_arith [*]) (Unsigned.cons r lsb)
  |]
termination_by w => w

-- Yet this desugared version works
def addU2 : {w x y : Nat} -> ⟪Unsigned w x⟫ -> ⟪Unsigned w y⟫ -> ⟪Bit c⟫ -> ⟪Unsigned w.succ (x + y + c)⟫
  | 0, 0, 0, _, _, cin => cast (by simp) [|Unsigned.cons .nil ~cin|]
  | w + 1, _, _, ux, uy, cin => Code.mk (stage := 0) $
    let .cons xbs xb := (Code.«~» (stage := 0) ux)
    let .cons ybs yb := (Code.«~» (stage := 0) uy)
    let .mk cin' lsb := bitAdd (Code.«~» (stage := 0) cin) xb yb
    let r := Code.«~» (stage := 0) $ addU2 [|xbs|] [|ybs|] [|cin'|]
    cast (by congr 1; simp_arith [*]) (Unsigned.cons r lsb)

[Dylan Ede]

I think it might be relevant that the desugared version (addU2) did not need termination_by w => w, and that without that line for addU, the error message is

fail to show termination for
  addU
with errors
structural recursion failed, produced type incorrect term
  unknown constant 'addU.match_2'

Is there any way of seeing what is being "produced" (in the sense of that error message) to see what is going wrong?

[Dylan Ede]

Bizarrely, I have found that if I place the definition for addU2 above addU, addU compiles without error.

[Kyle Miller]

By the way, rather than let stageLit := Lean.Syntax.mkNumLit $ toString stage you should be able to write let stageLit := quote stage or, later, just `(Code $(quote stage):num $ty)

[Kyle Miller]

Found the error: withEnv is way too strong, since it reverts the entire environment, not only your environment extension.

[Kyle Miller]

This is a fix:

def withStage [Monad m] [MonadFinally m] [MonadEnv m] (s : Nat) (x : m α) : m α := do
  let oldS ← getStage
  try
    setEnv <| stageEnv.setState (← getEnv) s
    x
  finally
    setEnv <| stageEnv.setState (← getEnv) oldS

[Kyle Miller]

I also modified the main file, but it doesn't seem necessary for the fix. The main thing is for ~ notation being sure to instantiate metavariables and having it make progress on instance synthesis problems before doing the hasFVar test. It's possible you also want to check that hasMVar is false since that's another way fvars can slip in, but in the future.

  -- The corresponding syntax in the toatie paper is ⟨_⟩
  syntax (name := «stx ⟪_⟫») "⟪" term "⟫" : term
  @[term_elab «stx ⟪_⟫»]
  def «elab ⟪_⟫» : TermElab
    | `(⟪ $ty:term ⟫), ty? => do
      let stage <- getStage
      withStage (stage + 1) <| elabTerm (<- `(Code $(quote stage):num $ty)) ty?
    | _, _ => throwUnsupportedSyntax

  -- The corresponding syntax in the toatie paper is ⟦_⟧
  syntax (name := «stx [|_|]»)"[|" term "|]" : term
  @[term_elab «stx [|_|]»]
  def «elab [|_|]» : TermElab
    | `([| $t |]), ty? => do
      let stage <- getStage
      withStage (stage + 1) <| elabTerm (<- `(Code.mk (stage := $(quote stage):num) $t)) ty?
    | _, _ => throwUnsupportedSyntax

  -- The corresponding syntax in the toatie paper is ˜_
  syntax:max (name := «stx ~_») "~" term:max : term
  @[term_elab «stx ~_»]
  def «elab ~_» : TermElab
    | `(~ $t), ty? => do
      let stage <- getStage
      let «⟪t⟫» <- withStage (stage - 1) <| withSynthesize <| elabTerm (<- `(Code.«~» (stage := $(quote (stage - 1)):num) $t)) ty?
      Term.synthesizeSyntheticMVars
      let «⟪t⟫» <- instantiateMVars «⟪t⟫»
      if stage == Nat.zero && «⟪t⟫».hasFVar then
        throwError "Free variables not allowed when escaping to stage 0"
      pure «⟪t⟫»
    | _, _ => throwUnsupportedSyntax

[Dylan Ede]

Kyle Miller said:

Found the error: withEnv is way too strong, since it reverts the entire environment, not only your environment extension.

Aha! D'oh. Seems obvious in hindsight! Thank you very much.

[Dylan Ede]

And thanks for the tip about free variables!

[Dylan Ede]

One thing I have noticed though, is that if I insert a tryPostpone near the end of «elab ~_», it appears that the correct stage environment is forgotten by the time the elaborator is attempted again - i.e. it thinks it's in stage 0 when it isn't. Any idea why that would be?

[Dylan Ede]

Hmm, looks like there's not much I can do about that because only what's in Lean.Elab.Term.SavedContext is preserved for the postponed elaborator call. So I guess I need to do something to prevent any postponement from within the elaboration of t, perhaps by catching postponement exceptions using withoutPostponing.

[Dylan Ede]

I found a solution to the postponing issue, which is to not use EnvExtension, and instead use Lean.Option, which is preserved in the SavedContext. The implementations of getStage and withStage then change to the following:

register_option hdl.stage : Nat := {
  defValue := 0
}

def getStage : TermElabM Nat := do
  return (← getOptions).getNat `hdl.stage

def withStage (stage : Nat) (x : TermElabM α) : TermElabM α := do
  withOptions (fun options => options.setNat `hdl.stage stage) x

The downside now is that a user could easily just set_option hdl.stage 1 and break the guarantees. Not sure if there's a way around that.

[Kyle Miller]

You don't have to register options, and in fact the defVal is not used by getNat.

If you don't register the option, users can't set_option, and furthermore you can make it very untypeable by using Name.num.

def hdlStageOptionName : Name := .num decl_name% 0

def getStage : Elab.Term.TermElabM Nat := do
  return (← getOptions).getNat hdlStageOptionName

[Dylan Ede]

Ah, thanks!