Cannot make st s a = s -> s & a as an instance of functor.

[katsujukou]

Sorry, but I have a question.

I'm trying to declare the functor instance for the state functor.
I have the state functor definition like this:

noeq
type st (s:Type) (a:Type) : Type = 
  | ST : k:(s -> s & a) -> st s a

And I have defined the functor typeclass like this:

noeq
type functor_laws
  (#f: Type -> Type)
  (#map : (#a:_ -> #b:_ -> (a -> b) -> f a -> f b))
  = {
    map_id : #a:_ -> (x:f a) -> 
      Lemma (map (id #a) x == id #(f a) x);
    map_compose: #a:_ -> #b:_ -> #c:_ -> g: (a -> b) -> h: (b -> c) -> x: (f a) -> 
      Lemma (map (fun (y: a) -> h (g y)) x == map h (map g x));
  } 

class functor (f: Type -> Type) = {
  map : (#a:Type) -> (#b:Type) -> (a -> b) -> f a -> f b;
  [@@@no_method]
  laws : unit -> GTot (functor_laws #f #map);
}

Now I'm attempting to define the functor instance as follows:

instance functor_st (#s: Type) : functor (st s)  
  = {
  map = (fun #a #b (f: a -> b) (stf: st s a) -> ST (fun s0 -> 
      let ST k = stf in
      let (s1, a1) = k s0 in
      (s1, f a1)
    )
  );

  laws = (fun _ -> {
    map_id = (fun _ -> ());
    map_compose = (fun _ _ _ -> ());
  });
}

But I got an error!

* Error 230 at src/Control.Monad.State.fst(26,40-26,41):
  - Variable "a#240" not found
  - See also src/Control.Monad.State.fst(26,14-26,15)

It seems like I have something wrong with the implicit argument #a,
but I cannot figure out how to correct it...

Additionally, instead of inlining the definition of map in the instance record,
I tried splitting it into a top-level function. With this change, it type-checked without above problems.

let 
  map_st #s #a #b (f: a -> b) (stf: st s a)
    : st s b 
    = ST (fun s0 -> 
      let ST k = stf in
      let (s1, a1) = k s0 in
      (s1, f a1)
    )

instance functor_st (#s: Type) : functor (st s)  
  = {
  map = map_st #s;

  laws = (fun _ -> {
    map_id = (fun _ -> ());
    map_compose = (fun _ _ _ -> ());
  });
}

(Though this time I encountered the error with laws part, stating that SMT solver could not proof automatically)

[nikswamy]

The variable not found is a bug; thanks for the bug report!

Here's a way to make your code work. Aside from using your top-level definition as a workaround for the scoping bug, the main issue here is that you need to work with FStar.FunctionalExtensionality. (https://fstar-lang.org/tutorial/book/part2/part2_equality.html#functional-extensionality)

open FStar.Tactics.Typeclasses
open FStar.FunctionalExtensionality
module F = FStar.FunctionalExtensionality

noeq
type functor_laws
  (#f: Type -> Type)
  (#map : (#a:_ -> #b:_ -> (a -> b) -> f a -> f b))
  = {
    map_id : #a:_ -> (x:f a) -> 
      Lemma (map (id #a) x == id #(f a) x);
    map_compose: #a:_ -> #b:_ -> #c:_ -> g: (a -> b) -> h: (b -> c) -> x: (f a) -> 
      Lemma (map (fun (y: a) -> h (g y)) x == map h (map g x));
  } 

class functor (f: Type -> Type) = {
  map : (#a:Type) -> (#b:Type) -> (a -> b) -> f a -> f b;
  [@@@no_method]
  laws : unit -> GTot (functor_laws #f #map);
}

// The ^-> notation from FStar.FunctionalExtensionality
// defines the type of functions for which extensionality is valid
let st (s:Type) (a:Type) : Type = (s ^-> s & a)

let map_st #s #a #b (f: a -> b) (stf: st s a)
: st s b 
= F.on_dom s //F.on_dom introduces a ^-> function type
    (fun s0 -> 
      let k = stf in
      let (s1, a1) = k s0 in
      (s1, f a1))

//These lemmas are valid for feq; a stepping stone to ==
let map_st_id (s:Type): 
    #a:_ -> (x:st s a) ->  Lemma (map_st (id #a) x `feq` id #(st s a) x)
  = fun #a x -> ()

let map_compose (s:Type):
  #a:_ -> #b:_ -> #c:_ -> g: (a -> b) -> h: (b -> c) -> x: (st s a) -> 
      Lemma (map_st (fun (y: a) -> h (g y)) x `feq` map_st h (map_st g x))
  = fun #a #b #c g h x -> ()


instance functor_st (#s: Type) : functor (st s)  
  = {
  map = map_st #s;

  laws = (fun _ -> {
    map_id = (fun #a x -> map_st_id s #a x);
    map_compose = (fun #a #b #c g h x -> map_compose s g h x);
  });
  }

[katsujukou]

Thank you for your response and for providing a detailed solution!
I understand the point.
By the way, I noticed that FStar.FunctionalExtensionality.on could also be a possible choice.
This is because our let st s a = s -> s & a is not a dependent arrow.
In the more general case where the return of the arrow depends on the argument value, on_domain would be the appropriate solution. Is this correct?