LazyList.ts

type Lazy<A> = () => A

function lazy<A>(fn: () => A): Lazy<A> {
    type State
        = { evaluated: false, evaluator: () => A }
        | { evaluated: true, value: A } 
    let state: State = { evaluated: false, evaluator: fn }
    return () => {
        if (state.evaluated === true) {
            return state.value
        } else {
            const value = state.evaluator()
            state = { evaluated: true, value }
            return value
        }
    }
}

function strict<A>(value: A): Lazy<A> {
    return () => value
}

type Nil = null
type LazyCons<A> = { head: A, tail: LazyList<A> }
type LazyList<A> = Lazy<LazyCons<A> | Nil>

function lazyList<A>(lst: A[]): LazyList<A> {
    return lazy(() => {
        if (lst.length <= 0) {
            return null
        } else {
            const [head, ...tail] = lst
            return { head, tail: lazy(lazyList(tail)) }
        }
    })
}

function generate<A>(fn: (prev: A) => A, initial: A): LazyList<A> {
    return lazy(() => {
        return {
            head: initial,
            tail: generate(fn, fn(initial))
        }
    })
}

function head<A>(lst: LazyList<A>): A | undefined {
    const evaluated = lst()
    if (evaluated === null) {
        return undefined
    } else {
        return evaluated.head
    }
}

function concat<A>(left: LazyList<A>, right: LazyList<A>): LazyList<A> {
    return lazy(() => {
        const evaluated = left()
        if (evaluated === null) {
            return right()
        } else {
            return {
                head: evaluated.head,
                tail: concat(evaluated.tail, right)
            }
        }
    })
}

function flatten<A>(lst: LazyList<LazyList<A>>): LazyList<A> {
    return lazy(() => {
        const evaluated = lst()
        if (evaluated === null) {
            return null
        } else {
            return concat(evaluated.head, flatten(evaluated.tail))()
        }
    })
}

function map<A, B>(lst: LazyList<A>, fn: (a: A) => B): LazyList<B> {
    return lazy(() => {
        const evaluated = lst()
        if (evaluated === null) {
            return null
        } else {
            return {
                head: fn(evaluated.head),
                tail: map(evaluated.tail, fn)
            }
        }
    })
}

function flatMap<A, B>(lst: LazyList<A>, fn: (a: A) => LazyList<B>): LazyList<B> {
    return flatten(map(lst, fn))
}

console.log(head(map(generate(i => i + 1, 0), i => {
    console.log(i)
    return i + '!'
})))