-- streams without memoization

type stream {
  cons : integer -> (unit -> stream) -> stream;
}

function force(f : unit -> stream) : stream {
  push unit;
  push f;
  tcall $ 1;
}

function filter(f : integer -> bool, s : unit -> stream, unit) : stream {
  push s;
  call force;
  tsave s1 {
    push s1.cons[0];
    push f;
    call $ 1;
    br {
      true: {
        push s1.cons[1];
        push f;
        calloc filter 2;
        push s1.cons[0];
        alloc cons;
        ret;
      }
      false: {
        push unit;
        push s1.cons[1];
        push f;
        tcall filter;
      }
    };
  };
}

function nth(n : integer, s : unit -> stream) : integer {
  push s;
  call force;
  tsave s1 {
    push n;
    push 0;
    eq;
    br {
      true: { push s1.cons[0]; ret; }
      false: {
        push s1.cons[1];
        push 1;
        push n;
        sub;
        tcall nth;
      }
    };
  };
}

function numbers(n : integer, unit) : stream {
  push n;
  push 1;
  add;
  calloc numbers 1;
  push n;
  alloc cons;
  ret;
}

function indivisible(n : integer, m : integer) : bool {
  push n;
  push m;
  mod;
  push 0;
  eq;
  br {
    true: { push false; ret; }
    false: { push true; ret; }
  };
}

function eratostenes(s : unit -> stream, unit) : stream {
  push s;
  call force;
  tsave s1 {
    push s1.cons[1];
    push s1.cons[0];
    calloc indivisible 1;
    calloc filter 2;
    calloc eratostenes 1;
    push s1.cons[0];
    alloc cons;
    ret;
  };
}

function primes() : unit -> stream {
  push 2;
  calloc numbers 1;
  calloc eratostenes 1;
  ret;
}

function main() {
  call primes;
  push 10;
  call nth;
  trace;
  pop;
  call primes;
  push 50;
  call nth;
  trace;
  pop;
  push void;
  ret;
}