runtime_analysis_with_the_typechecker.ts

// Let's use type annotations to examing the runtime/size of things!
// You'll have more fun if you can hover over variables and see the types for yourself.
// Copy paste into here: https://www.typescriptlang.org/play

interface SConstant { __constant_size: true }
interface SLinear<T extends string> { __variable_size: T }
interface SPlus<T extends Size, U extends Size>{ __splusand_1: T, __splusand_2: U }
interface STimes<T extends Size, U extends Size>{ __stimesand_1: T, __stimesand_2: U }
type Size = SConstant | SLinear<string> | SPlus<Size, Size> | STimes <Size, Size>
// Declare some type brands, and a union of all the possibilites 


type SizeMonad<T, S extends Size> = T & S
// We can declare a "Monad" that only exists within the type system, to glue this value onto real values

// Simple `return` function for wrapping
function returnSizeMonad<T, S extends Size = SConstant>(a: T): SizeMonad<T, S > { 
    return a as unknown as T & S
}

// Declare some arrays with named linear sizes
const big = returnSizeMonad<number[], SLinear<"sizeof(big)">>([1,2,3,4,5,6,7,8,9])
const sm = returnSizeMonad<number[], SLinear<"sizeof(sm)">>([1,2,3])

// Lets build a concatenator that handles the types properly
function concat<SX extends Size, SY extends Size, T>(
    xs: SizeMonad<T[], SX>, 
    ys: SizeMonad<T[], SY>, 
    ): SizeMonad<T[], SPlus<SX,SY>> {
        const res: T[] = [...xs, ...ys]
    return res as SizeMonad<T[], SPlus<SX,SY>>
}

const concatenated = concat(big,sm)
//    concatenated: SizeMonad<number[], SPlus<SLinear<"sizeof(big)">, SLinear<"sizeof(sm)">>>;
// Neato: the types know how big it is!



// Lets define a similar family of types and monad for the runtime to calculate a value, using these sizes

interface Constant { __constant_time: true }
interface Linear<T extends string> { __variable_time: T }
interface Plus<T extends Runtime, U extends Runtime>{ __plusand_1: T, __plusand_2: U }
interface Times<T extends Runtime, U extends Runtime>{ __timesand_1: T, __timesand_2: U }
type Runtime = Constant | Linear<string> | Plus<Runtime, Runtime> | Times <Runtime, Runtime> | Size

type RuntimeMonad<T, R extends Runtime> = T & R

// Consider a simple array sum function, it's runtime is exactly proportional to the size of the array being summed
function sumOf<S extends Size>(xs: SizeMonad<number[], S>): RuntimeMonad<number, S> {
    let s = 0
    xs.forEach((x) => { s += x })
    return s as number & S
}

const sumBig = sumOf(big)
//    sumBig: RuntimeMonad<number, SLinear<"sizeof(big)">>;
const sumSm = sumOf(sm)
//    sumSm: RuntimeMonad<number, SLinear<"sizeof(sm)">>;
// Obvious stuff, but it's a foundation

// Lets build a way to combines two calculations with RuntimeMonads attached
function withBoth<TT extends Runtime, UT extends Runtime, T, U, R>(
    t: RuntimeMonad<T, TT>, 
    u: RuntimeMonad<U, UT>, 
    f: (t: T, u: U) => R
    ): RuntimeMonad<R, Plus<TT,UT>> {
        const res: R = f(t, u)
    return res as RuntimeMonad<R, Plus<TT,UT>>
}

// Now we can sum our sums, and seen the runtime is also the sum 
const sumBoth1 = withBoth(sumBig,sumSm, (a,b) => a + b )
//    sumBoth1: RuntimeMonad<number, Plus<SLinear<"sizeof(big)">, SLinear<"sizeof(sm)">>>;

// We can also show this is equivalient to summing the concatenated array!
const sumBoth2 = sumOf(concatenated)
//    sumBoth2: RuntimeMonad<number, SPlus<SLinear<"sizeof(big)">, SLinear<"sizeof(sm)">>>;

// You can also consider a `then` for chaining operations
function then<A,B, RTA extends Runtime, RTB extends Runtime>(a: RuntimeMonad<A, RTA>, f: (a: A) => RuntimeMonad<B, RTB>): RuntimeMonad<B, Plus<RTA,RTB>>{
    return f(a) as unknown as RuntimeMonad<B, Plus<RTA,RTB>>
}

// Maps are where all the fun is though, let's take a sized array, a transformation function that
// returns a RuntimeMonad AND a SizeMonad, and the bucket of type parameters we need to make it all work
// ultimately producing a map that preserves our size and runtime calculations
function arrayMap<A, B, SA extends Size, RTF extends Runtime = Constant, SR extends Size = SConstant >(
    a: SizeMonad<A[], SA>, 
    f: (a: A) => RuntimeMonad<SizeMonad<B, SR>, RTF>,
    ): RuntimeMonad<SizeMonad<B[], STimes<SA, SR>>, Times<RTF, SA>> {
        return a.map(f) as unknown as RuntimeMonad<SizeMonad<B[], STimes<SA, SR>>, Times<RTF, SA>>
}


// A practical application is determining the size and runtime of a cross product calculation
function crossProduct<
        SA extends Size, 
        SB extends Size
    >(as: SizeMonad<number[], SA>, bs: SizeMonad<number[], SB>)
{   
    // Define a row multiplication function using arrayMap
    function scaleBsBy(a: number){
        const multiplyRowByA = (b: number): number & Constant & SConstant => (a * b as unknown as number & Constant & SConstant)
        // Some jiggry pokery is required to indicate the row multiplication uses constant time and space per element in B
        const res = arrayMap<number, number, SB, Constant, SConstant>(bs, multiplyRowByA)
        return res
    }
    // The real fun part, we can use our arrau map, with the scaleBsBy function, 
    // which is built with an array map, and all the types work out
    const res = arrayMap(as, scaleBsBy)
    return res
}

const crossed = crossProduct(big, sm)
// declare const crossed: RuntimeMonad<
//          SizeMonad<number[][], 
//          STimes<SLinear<"sizeof(big)"> & Constant, STimes<SLinear<"sizeof(sm)"> & Constant, SConstant>>>, 
//      Times<Times<Constant, SLinear<"sizeof(sm)"> & Constant>, SLinear<"sizeof(big)"> & Constant>>;
//
// The Runtime and size of crossed both reduce to sizeof(big) * sizeof(sm) neat!