# Polymorphic Sorts

Did we ever have a discussion about adding `Set` and `Map`
as indices? If so what was the conclusion? I think we’d need some sort (ha) of polymorphism to allow for sets of bool/int/etc?

We need two things:

1. Support polymorphism in sort checking if we want to have polymorphic functions on sets, e.g.,

    union<s>(Set<S>, Set<S>) -> Set<S>

The sort checking algorithm is half bidirectionalish, half hindley-milnerish, with the added complexity that it supports some implicit coercions (which we need to `elaborate` before `conv`).

Adding polymorphism to it would be fun... Maybe adding explicit ascriptions can help.

Presumably we want to use polymorphic sorts on polymorphic types and instantiate the polymorphic sort with the sort associated to the type parameter. We can only do that if the type parameter is of kind base. We currently don't support type parameters of kind base in adts. I don't see any issue with supporting them though.

The one pesky problem is that because we have abstract refinements as indices, not all sorts support equality, so it would be unsound to naively have a `Set<S>` for any `s`

I think the most general way to solve this is to have sort bounds. Say for example, you want to index a `HashMap<K, V>` by the set of keys, you could explicitly say that (the sort associated to) `K` needs to support equality

#[flux::refined(keys: Set<sortof(K)> where sortof(T) has Eq)]
struct HashMap<K, V> {...}

(I'm not suggesting the sortof syntax, just trying to make explicit that the bound is on the sort associated to the type parameter)
To alleviate this we could make has `Eq` the default and if you want to lift it you could use hash `?Eq` similar to how rust has `?Sized` to lift the sized restriction on type parameters.

Some other options to solve this include

1. Separate abstract refinement syntactically from other indices so we have equality for all indices. The problem here is to find syntax for it because you will have types that can be parameterized by types, abstract refinements, and indices. That's a lot of parentheses `(MyType<t>(|x| x > 0|)[10] ?)`. Maybe there's some other way to have syntax for this.

2. Transition away from indices as fields of a transparent record sort (with structural equality) into indices as measures (uif) over some opaque sort (that always has equality). I don't know about this. I think both make sense. If you are capturing the semantics of a type exactly, then it makes a lot of sense to have structural equality for its sort, e.g.,

#[flux::refined_by(a: int, b: int)]
struct Pair {
    x: i32[a],
    y: i32[b],
}

Although, the cases where you exactly capture the semantics of a type are probably rare. There may be space to have both opaque and transparent sorts.

## A Monomorphic Map

/// define a type indexed by a map
#[flux::opaque]
#[flux::refined_by(map: Map<int, int>)]
pub struct RMap {
    inner: std::collections::HashMap<i32, i32>,
}

impl RMap {
    #[flux::trusted]
    pub fn new() -> Self {
        Self { inner: std::collections::HashMap::new() }
    }

    #[flux::trusted]
    #[flux::sig(fn(self: &strg RMap[@m], k: i32, v: i32) ensures self: RMap[map_set(m, k, v)])]
    pub fn set(&mut self, k: i32, v: i32) {
        self.inner.insert(k, v);
    }

    #[flux::trusted]
    #[flux::sig(fn(&RMap[@m], k: i32) -> Option<i32[map_get(m, k)]>)]
    pub fn get(&self, k: i32) -> Option<i32> {
        self.inner.get(&k).copied()
    }
}

## A Polymorphic Map

Key idea: whenever you see a rust type-variable appear as a `sort` we mean

* externally, the *sort-of-the-thing-that-`K`-is-instantiated-with* at instantiation sites, or
* internally, an opaque `K` sort when `K` is quantified.

/// define a type indexed by a map
#[flux::opaque]
#[flux::refined_by(map: Map<K, V>)]
pub struct RMap<K, V> {
    inner: std::collections::HashMap<K, V>,
}

impl RMap<K, V> {
    #[flux::trusted]
    pub fn new() -> Self {
        Self { inner: std::collections::HashMap::new() }
    }

    #[flux::trusted]
    #[flux::sig(fn<K as base, V as base>(self: &strg RMap<K,V>[@m], k: K, v: V) ensures self: RMap[map_set(m, k, v)])]
    pub fn set(&mut self, k: &K, v: i32) {
        self.inner.insert(k, v);
    }

    #[flux::trusted]
    #[flux::sig(fn(&RMap[@m], k: &K) -> Option<&V[map_get(m, k)]>)]
    pub fn get(&self, k: i32) -> Option<i32> {
        self.inner.get(&k).copied()
    }
}

Here's the LH version https://github.com/ucsd-progsys/liquidhaskell/blob/c55c331fc1d4724c026f144ce44759609b9927ad/tests/pos/Maps.hs#L29

One **advantage** over the LH version would be something like this:

#[flux::refined_by(a: int, b: int)]
struct Pair {
    x: i32[a],
    y: i32[b],
}

fn test() {
    let p0 = Pair {x: 10, y:10 };
    let p1 = Pair {x: 20, y:20 };
    
    let mut m = RMap::new();
    m.set(10, p0);
    m.set(20, p1);
   
    // LH would *fail* to verify the below (without abstract-refinements on `RMap`)
    assert(m.get(10).unwrap().x == 10);
    assert(m.get(20).unwrap().x == 20);
}

## RSet

#![flux::defs {
     fn set_add<T>(x: int, s: Set<T>) -> Set<T> { set_union(set_singleton(x), s) }
}]

use std::hash::Hash;

#[flux::opaque]
#[flux::refined_by(<T> { elems: Set<T> } )]        implicit or explicit
pub struct RSet<T> {
    inner: std::collections::HashSet<T>,
}

RSet<nat>

RSet<\v. {w. (i32[w], i32[v]) | v >= 0}>



RSet<λv. {i32[v] | v >= 0}>


RSet<λv. {w. (i32[w], i32[v]) | v >= 0}>

#[flux::sig(fn<T as base, refine s: Set<T>>(set: &strg RSet<T>[s], elem: T) ensures set: RSet[ set_union(set_singleton(x), s)])]
pub fn insert<T>(set: &mut RSet<T>, elem: T) {
   set.inner.insert(elem);
}

// RSet<Nat> 

impl<T: Eq + PartialEq + Hash> RSet<T> {
    #[flux::trusted]
    #[flux::sig(fn() -> RSet[set_empty(0)])]
    pub fn new() -> Self {
        Self { inner: std::collections::HashSet::new() }
    }

    // ensures self: RSet[set_add(k, s.elems)])]
    #[flux::trusted]
    #[flux::sig(fn(self: &strg RSet[@s], elem: T)
                ensures self: RSet[ set_union(set_singleton(x), s) ])]
    pub fn insert(&mut self, elem: T) {
        self.inner.insert(elem);
    }

    #[flux::trusted]
    #[flux::sig(fn(&Set[@s], &T[@elem]) -> bool[set_is_in(elem, s.elems)])]
    pub fn contains(&self, elem: &T) -> bool {
        self.inner.contains(elem)
    }
}

NL: Look at signature for `insert` -- what can you
*instantiate* the `T` with?

ANS: anything **with indices** -- this is most things,
but not type-alias-projection-mumble which have NO INDEX.

BASE = Stuff that can have an index -- i.e. things in `BaseTy`

data Set a = Leaf | Node (Set a) (Set a)

fn foo(bob: &mut RSet<Nat>) {
    bob.insert(10)
}

fn foo(b: &mut RVec<Nat>) {
    b.push(10)
}

fn insert<T as base>(self: &mut RSet<T>[@s], elem: T) ensures self: RSet<T>[set_add(s, elem)] 


// This is what happens now during desugaring
fn insert<T as base>(self: &mut RSet<∃v. T[v]>[@s], elem: T) ensures self: RSet<T>[set_add(s, elem)] 


fn insert<T as base, s: Set<T>>(self: &mut RSet<T>[s], elem: T) 
   ensures self: RSet<T>[set_add(s, elem)] 




#[flux::sig(fn<T as base>(b: bool, x: T[@n], y: T[@m]) -> T[if b { n } else { m }])]
fn choose<T>(b: bool, x: T, y: T) -> T {
    if b {
        x
    } else {
        y
    }
}

fn test01() {
    assert(choose(false, 0, 1) == 0);
    assert(choose(true, 0, 1) == 1);
}

// choose<T := λx. {i32[x] | $k(x)}, n := 0, m := 1>(false, 0 ,1)
// choose : fn(bool, (λx. {i32[x] | $k(x)})[0], (λx. {i32[x] | $k(x)})[1])
// choose : fn(bool, {i32[0] | $k(0)}, {i32[1] | $k(1)})

https://github.com/flux-rs/flux/blob/64612d726a69f28c9668506f04580e05bd03dc00/crates/flux-tests/tests/neg/surface/refined_type_var00.rs#L10

Current issue:

#[flux::opaque]
#[flux::refined_by(elems: Set<Tiger>)]
pub struct RSet<Tiger> {
    pub inner: std::collections::HashSet<Tiger>,
}

#[flux::sig(fn<A as base>(set: &RSet<A>[@s], &A[@elem]) -> bool[set_is_in(elem, s.elems)])] 

s       : RSet<A> 
s.elems : Set<A>
    

sort error because:

s : RSet
.elems : Set<T>

and so

s.elems: Set<T>

which clashes with the fact that

elem: A

and so we cannot check

set_is_in(elem, s.elems)

Instead we need to

0. Track the `SortGenerics` for each ADT/DefId
- RSet -> [Tiger]
- Booger -> [A]
- For now, use _all_ the generics, i.e. just `generics_of`

1. Change the type `Record` to track the arguments for `SortGenerics`
- fhir::Sort::Record(DefId, List<Sort>)

2. Change the field lookup `e.fld` so that
synth_expr(e.fld) = T[A... := args...]
where
synth_expr(e) ==> Record(def_id, args...)
SortGenerics(def_id) ==> A...
synth_field(def_id, fld) => T

e.g.

synth_expr(s.elems) ==> Set<T> [ T := i32 ] ==> Set<int>
as
synth_expr(e) ==> Record(RSet, [int])
SortGenerics(RSet) ==> [T]
synth_field(RSet.elems) ==> Set<T>

#[flux::sort_generics(A as base)]
#[flux::refined_by(elems:Set<A>)]
struct Booger<A, B> { 
    inner: Set<A>
}

#[flux::opaque]
#[flux::refined_by(elems: Set<sort_of_indices_of(T)>)]
pub struct RSet<T> {
    pub inner: std::collections::HashSet<T>,
}

    
#[flux::sig(fn(RSet<i32{v: v > 0}>[@n]) )]
fn bob(x: RSet<i32>) {
  ...    
}
    
    
    

#[flux::sig(fn<A as base>(x:A) -> RSet<A>[set_singleton(x)])]
fn bob<A>(x:A) -> RSet<A> {
  ...    
}
  
fn test() {
    let x:i32 = 10;
    let s = bob(x);
    assert(s.contains(&x))
}

e.g.

Subst T := i32[10] in Set<sort_of(T)>

    T := i32[10] 
    -->    
    Set<sort_of(i32[10])>
    -->
    Set<sort_of(i32)>
    -->
    Set<int>

Sort =
    | Int
    | Bool
    | Param(DefId)

## Trying to break it

    
use std::hash::Hash;

mod test01 {
    // This test is, fine just explaining why it works.

    use super::*;

    #[flux::sig(fn<T as base>(s: RSet<T>, x: T{set_is_in(x, s.elems)}))]
    fn foo<T: Eq + Hash>(s: RSet<T>, x: T) {}

    #[flux::sig(fn(s: RSet<i32{v: v > 0}>{set_is_in(1, s.elems)}))]
    fn client(s: RSet<i32>) {
        foo(s, 1);
    }

    // The type of `foo` is
    // for<T: base, s: Set<T::sort>, x: T::sort>
    //     fn(RSet<∃v: T::sort. T[v]>[s], { T[x] | set_is_in(x, s) })
    //
    // The type of `client` is
    // for<s': Set<int>>
    //     fn({ RSet<∃v: int. {i32[v] | v > 0}>[s'] | set_is_in(1, s') })
    //
    // Note that the `T` in `RSet<T>` is always something in `TyKind`.
    //
    // When calling `fun` we instantiate it like so
    // fun<T:=λv:int {i32[v] | κ(v)}, s:=s', x:=1>
    //      => fn(RSet<∃v: int. {i32[v] | κ(v)}>[s'], {i32[1] | set_is_in(1, s')})
    //
    // This then generate subtyping constraints
    // * RSet<∃v: int. {i32[v] | v > 0}>[s'] <: RSet<∃v: int. {i32[v] | κ(v)}>[s']
    // * i32[1] <: {i32[1] | set_is_in(1, s')}
    //
    // In context where `s': RSet<int>, set_is_in(1, s')`
}

mod test02 {
    use super::*;

    /// TODO: This currently crashes. We should reject it because `RSet<impl Eq + Hash>` 
    /// is not well-formed, because `T` in `RSet` is "special" (does not admit `impl`)
    fn foo() -> RSet<impl Eq + Hash> {
        RSet::<i32>::new()
    }

    /// TODO: This currently crashes. We should accept it. Nothing wrong with using
    /// an impl trait with Option. passes in `main` but BREAKS in `rset` branch; so 
    /// should allow it in `rset`
    fn baz() -> Option<impl Eq + Hash> {
        Some(1)
    }
}
    
fn mk_eq_hash() -> impl Eq + Hash {
    0
}
    
mod test03 {
    use super::*;

    /// TODO: This currently crashes. We should gracefully reject it.
    fn foo() {
        let x = mk_eq_hash(); // x: impl Eq + Hash
        // This will try to create an `RSet<impl Eq + Hash>` which is not 
        // well-formed. We should reject it with a *kind mismatch error*. 
        // Note that this is an error we can only catch during refinement
        // type checking.
        let mut s = RSet::new();
        s.insert(x);
    }

    /// We currently accept this function, but it should fail because `RSet<T>` is 
    /// not well-formed if `T` is of kind type. I don't know if this leads to 
    /// unsoundness, but you can create ill-formed types instantiating the parameter 
    /// with something that doesn't have a sort like so:
    ///
    /// ```ignore
    /// trait SomeTrait { type Assoc: Eq + Hash }
    /// let f = baz::<<i32 as SomeTrait>::Assoc>;
    /// ```
    /// However, writing the above requires function pointers which we don't currently supported
    #[flux::sig(fn(RSet<T>[@s]))]
    fn baz<T: Eq + Hash>(set: RSet<T>) {}
}

mod test04 {
    use super::*;

    #[flux::sig(fn<T as base>(x: T))]
    fn foo<T>(x: T) {}

    // TODO: variant of `test3::foo`
    fn client() {
        let x = mk_eq_hash();
        // Fails with parameter inference call but should fail with kind error. This is unrelated
        // to the changes to sort.
        foo(x);
    }
}



#[flux::opaque]
#[flux::refined_by(elems: Set<T>)]
pub struct RSet<T> {
    inner: std::collections::HashSet<T>,
}

#[flux::generics(T as base)]
impl<T> RSet<T> {
    #[flux::trusted]
    #[flux::sig(fn() -> RSet<T>[set_empty(0)])]
    pub fn new() -> RSet<T> {
        let inner = std::collections::HashSet::new();
        RSet { inner }
    }

    #[flux::trusted]
    #[flux::sig(fn(set: &strg RSet<T>[@s], elem: T) ensures set: RSet<T>[set_union(set_singleton(elem), s)])]
    pub fn insert(self: &mut Self, elem: T)
    where
        T: Eq + Hash,
    {
        self.inner.insert(elem);
    }

    #[flux::trusted]
    #[flux::sig(fn(set: &RSet<T>[@s], &T[@elem]) -> bool[set_is_in(elem, s.elems)])]
    pub fn contains(self: &Self, elem: &T) -> bool
    where
        T: Eq + Hash,
    {
        self.inner.contains(elem)
    }
}

#[flux::sig(fn(RSet<i32{v: v > 0}>[@s], y:i32{set_is_in(y, s.elems)}) )]
pub fn test2(s: RSet<i32>, y: i32) {
    assert(contains(&s, &y));
}

#[flux::sig(fn<T as base>(RSet<∃v. T[v]>[@s], y:T{set_is_in(y, s.elems)}))]
pub fn test3<T>(s: RSet<T>, y: T)
where
    T: Eq + Hash,
{
    assert(contains(&s, &y));
}

#[flux::sig(fn(RSet<i32{v: 0 <= v}>[@s], y:i32{0 <= y && set_is_in(y, s.elems)}) )]
pub fn test4(s: RSet<i32>, y: i32) {
    test3(s, y)
}
    
// TODO: (see `baz` above) should be rejected, as `RSet<T>[@s]` is not WF UNLESS T is 'special' (as base => special but could be special otherwise too?)
#[flux::sig(fn(RSet<T>[@s]))]
pub fn test5<T>(s: RSet<T>)
where
    T: Eq + Hash,
{
}