Playing with streams as codata in Polarity (see: https://polarity-lang.github.io/oopsla24/)

streams.pol

-- ------------------------------- --
-- Non-dependent functions         --
-- ------------------------------- --

codata Fun(a b: Type) {
    Fun(a, b).ap(a b: Type, x: a) : b,
}

-- ------------------------------- --
-- Top type                        --
-- ------------------------------- --

-- This is useful for defining top-level definitions

data Top {
    Unit,
}

def Top.id(a: Type): a -> a {
    Unit => \x. x,
}

-- The above is sugar for:
--
--     def Top.id(a: Type): Fun(a, a) {
--         Unit => comatch Id {
--             ap(_, _, x) => x,
--         },
--     }

-- ------------------------------- --
-- Natural numbers                 --
-- ------------------------------- --

data Nat { 
  Z, 
  S(n : Nat),
}

def Nat.add(y: Nat): Nat {
    Z => y,
    S(x) => S(x.add(y)),
}

def Nat.mul(y: Nat): Nat {
    Z => Z,
    S(x) => x.add(y.mul(x)),
}

-- ------------------------------- --
-- Finite natural numbers          --
-- ------------------------------- --

data Fin(n: Nat) {
  FZ(n: Nat): Fin(S(n)),
  FS(n: Nat, i: Fin(n)): Fin(S(n)),
}

def Fin(n).to_nat(n: Nat): Nat {
    FZ(_) => Z,
    FS(n, i) => S(i.to_nat(n)),
}

def Fin(n).weaken(n: Nat): Fin(S(n)) {
    FZ(n) => FZ(S(n)),
    FS(n, i) => FS(S(n), i.weaken(n)),
}

def Fin(n).weaken_n(n: Nat, m: Nat): Fin(n.add(m)) {
    FZ(n) => FZ(n.add(m)),
    FS(n, i) => FS(n.add(m), i.weaken_n(n, m)),
}

-- ------------------------------- --
-- Lists                           --
-- ------------------------------- --

data List(a: Type) {
    Nil(a: Type): List(a),
    Cons(a: Type, x: a, xs: List(a)): List(a),
}

def List(a).append(a: Type, ys: List(a)): List(a) {
    Nil(_) => ys,
    Cons(_, x, xs) => Cons(a, x, xs.append(a, ys)),
}

-- ------------------------------- --
-- Monoids                         --
-- ------------------------------- --

codata Monoid(a: Type) {
    Monoid(a).mempty(a: Type): a,
    Monoid(a).mappend(a: Type, x: a, y: a): a,
}

codef AdditiveMonoid(a: Type): Monoid(Nat) {
    mempty(_) => 0,
    mappend(_, x, y) => x.add(y),
}

codef MultiplicativeMonoid(a: Type): Monoid(Nat) {
    mempty(_) => 1,
    mappend(_, x, y) => x.mul(y),
}

codef ListMonoid(a: Type): Monoid(List(a)) {
    mempty(_) => Nil(a),
    mappend(_, xs, ys) => xs.append(a, ys),
}

-- ------------------------------- --
-- Non-dependent product types     --
-- ------------------------------- --

-- | Negatively defined non-dependent product types
codata Prod(a b: Type) {
    -- | The first element of the pair
    Prod(a, b).fst(a b: Type): a,
    -- | The second element of the pair
    Prod(a, b).snd(a b: Type): b,
}

-- | Construct a pair from two elements
codef Pair(a b: Type, x: a, y: b): Prod(a, b) {
    fst(_, _) => x,
    snd(_, _) => y,
}

-- ------------------------------- --
-- Infinite streams                --
-- ------------------------------- --

-- Some operations were inspired by Agda’s guarded stream module:
-- https://agda.github.io/agda-stdlib/master/Codata.Guarded.Stream.html

codata Stream(a: Type) {
    -- | The first element of the stream
    Stream(a).head(a: Type): a,
    -- | The rest of the stream
    Stream(a).tail(a: Type): Stream(a),
}

-- | Add a new element to the start of the stream
codef SCons(a: Type, x: a, xs: Stream(a)): Stream(a) {
    head(_) => x,
    tail(_) => xs,
}

codef Repeat(a: Type, x: a): Stream(a) {
    head(_) => x,
    tail(_) => Repeat(a, x),
}

codef Map(a b: Type, xs: Stream(a), f: a -> b): Stream(b) {
    head(_) => f.ap(a, b, xs.head(a)),
    tail(_) => Map(a, b, xs.tail(a), f),
}

codef ScanL(a b: Type, xs: Stream(a), init: b, f: b -> a -> b): Stream(b) {
    head(_) => init,
    tail(_) => ScanL(a, b, xs.tail(a), f.ap(b, a -> b, init).ap(a, b, xs.head(a)), f),
}

codef Unfold(a b: Type, seed: a, next: a -> Prod(a, b)): Stream(b) {
    head(_) => next.ap(a, Prod(a, b), seed).snd(a, b),
    tail(_) => Unfold(a, b, next.ap(a, Prod(a, b), seed).fst(a, b), next),
}

codef Iterate(a: Type, init: a, f: a -> a): Stream(a) {
    head(_) => init,
    tail(_) => Iterate(a, f.ap(a, a, init), f),
}

codef Interleave(a: Type, xs: Stream(a), ys: Stream(a)): Stream(a) {
    head(_) => xs.head(a),
    tail(_) => Interleave(a, ys, xs.tail(a)),
}

codef Fib(n m : Nat): Stream(Nat) {
    head(_) => n,
    tail(_) => Fib(m, n.add(m)),
}

def Top.ones: Stream(Nat) {
    Unit => Repeat(Nat, 1),
}

def Top.nats_from(init: Nat) : Stream(Nat) {
    Unit => Iterate(Nat, init, \n. S(n)),
}

def Top.nats : Stream(Nat) {
    Unit => Unit.nats_from(0),
}

def Top.tabulate_from(a: Type, init: Nat, f: Nat -> a) : Stream(a) {
    Unit => Map(Nat, a, Unit.nats_from(init), f)
}

def Top.tabulate(a: Type, f: Nat -> a) : Stream(a) {
    Unit => Unit.tabulate_from(a, 0, f)
}

def Top.take(a: Type, xs: Stream(a), n: Nat): List(a) {
    Unit => n.match {
        Z => Nil(a),
        S(n) => Cons(a, xs.head(a), Unit.take(a, xs.tail(a), n)),
    }
}

def Top.drop(a: Type, xs: Stream(a), n: Nat): Stream(a) {
    Unit => n.match {
        Z => xs,
        S(n) => Unit.drop(a, xs, n).tail(a),
    }
}

def Top.nth(a: Type, xs: Stream(a), n: Nat): a {
    Unit => Unit.drop(a, xs, n).head(a)
}

def List(a).append_stream(a: Type, ys: Stream(a)): Stream(a) {
    Nil(_) => ys,
    Cons(_, x, xs) => SCons(a, x, xs.append_stream(a, ys)),
}

-- ------------------------------- --
-- Tests                           --
-- ------------------------------- --

data Eq(a: Type, x y: a) {
    Refl(a: Type, x: a): Eq(a, x, x),
}

def Top.example_three_nats: List(Nat) {
    Unit => Cons(Nat, 0, Cons(Nat, 1, Cons(Nat, 2, Nil(Nat)))),
}

def Top.test_take_nats: Eq(List(Nat), Unit.take(Nat, Unit.nats, 3), Unit.example_three_nats) {
    Unit => Refl(List(Nat), Unit.example_three_nats)
}

def Top.test_nth_nats: Eq(Nat, Unit.nth(Nat, Unit.nats, 3), 3) {
    Unit => Refl(Nat, 3)
}

def Top.test_drop_head: Eq(Nat, Unit.drop(Nat, Unit.nats, 5).head(Nat), 5) {
    Unit => Refl(Nat, 5),
}