Exploring building a CSS-in-JS system in Typescript

style-props.tsx

import { CSSProperties, Component } from "react"

/**
 * Base types.
 */

/**
 * A SheetProducer makes a sheet based on a context.
 * Right now, the context only contains a `theme`.
 */
type SheetProducer<Theme, Props, Classes extends string> = (
	context: Context<Theme>
) => Sheet<Props, Classes>

/**
 * A sheet contains maps names to styles.
 * Each style is for a single class of elements.
 */
type Sheet<Props, Classes extends string> = {
	[K in Classes]: Styles<Props>
}

// TODO: assert StaticSheet is assignable to Sheet<never>

/**
 * A static stylesheet that can be used in a React component's render() function.
 */
type StaticSheet<S extends Sheet<any, any>> = {
	[K in keyof S]: CSSProperties
}

type Styles<Props> = {
	[K in keyof CSSProperties]:
		| CSSProperties[K] // A static property for this class
		| ((props: Props) => CSSProperties[K]) // A dynamic property based on props. style?
} & {
	// Bulk dynamic properties. style?
	// This is a stretch feature - it might not be worth including.
	_?: (props: Props) => CSSProperties
}

type Context<Theme> = { theme: Theme }

/**
 * StylesCache memoizes the computation of `StaticSheet`s by its input parameters.
 * The StylesCache is the core of our engine.
 */
class StylesCache<Theme extends object> {
	private byProducer: WeakMap<
		SheetProducer<Theme, any, any>,
		{
			byContext: WeakMap<
				Context<Theme>,
				{
					sheet: Sheet<any, any>
					byProps: WeakMap<
						object, // Some props
						{
							staticSheet: StaticSheet<any>
						}
					>
				}
			>
		}
	> = new WeakMap()

	/**
	 * Get a static sheet suitable for rendering a component.
	 */
	getStaticSheet<Props extends object>(
		producer: SheetProducer<Theme, Props, any>,
		context: Context<Theme>,
		props: Props
	): StaticSheet<any> {
		let producerCache = this.byProducer.get(producer)
		if (!producerCache) {
			producerCache = { byContext: new WeakMap() }
			this.byProducer.set(producer, producerCache)
		}

		let contextCache = producerCache.byContext.get(context)
		if (!contextCache) {
			contextCache = {
				sheet: producer(context),
				byProps: new WeakMap(),
			}
			producerCache.byContext.set(context, contextCache)
		}

		let propsCache = contextCache.byProps.get(props)
		if (!propsCache) {
			propsCache = {
				staticSheet: getStaticSheetFromSheetAndProps(contextCache.sheet, props),
			}
			contextCache.byProps.set(props, propsCache)
		}

		// Further optimizations left on the table:
		// - Memoize extracting the "static parts" from the sheet
		// - Memoize detecting the "dynamic" keys inside the sheet
		// - Considering deepEquals or shallowEquals, etc. Right now only referrential.

		return propsCache.staticSheet
	}
}

function getStaticSheetFromSheetAndProps<Props, S extends Sheet<Props>>(
	sheet: S,
	props: Props
): StaticSheet<S> {
	const result: StaticSheet<S> = {} as any

	for (const [className, styles] of exactEntries(sheet)) {
		const { _: getOverrideProperties, ...styles2 } = styles
		const resultStyles: any = {} // XXX

		// Props by key
		for (const pair of exactEntries(styles2)) {
			if (!pair) {
				continue
			}
			const [property, propertyValue] = pair
			if (typeof propertyValue === "function") {
				const computedPropertyValue = propertyValue(props)
				resultStyles[property] = computedPropertyValue as any // XXX
			} else {
				resultStyles[property] = propertyValue as any // XXX
			}
		}

		if (getOverrideProperties) {
			Object.assign(resultStyles, getOverrideProperties(props))
		}

		result[className] = resultStyles
	}

	return result
}

type ExactEntries<T> = {
	[K in keyof T]: [K, T[K]]
}[keyof T]

/**
 * Variant of Object.entries that assumes T is an exact object.
 * Object.entries is pessimistically typed.
 * @see https://github.com/microsoft/TypeScript/pull/12253#issuecomment-263132208
 */
function exactEntries<T>(obj: T): Array<ExactEntries<T>> {
	// Using Array due to not setting --downlevelIteration when targeting whateverthefuck.
	// TODO: change to just Object.entries, or use lodash or something.
	return Array.from(Object.entries(obj)) as any
}

/**
 * StyledComponent
 *
 * oof, this needs a lot of work to be ergonomic.
 */
type MyTheme = {
	name: "dark" | "light"
}
const globalStylesCache = new StylesCache<MyTheme>()
const globalContext: Context<MyTheme> = {
	theme: { name: "dark" },
}

// Take 1: an Abstract Class.
// Learnings: abstract classes are obnoxious, and defining methods on the subclass does
// not infer/fill type variables.
abstract class StyledComponent<
	Props,
	SP extends SheetProducer<MyTheme, Props, any> = SheetProducer<
		MyTheme,
		Props,
		any
	>
> extends Component<Props> {
	// Subclasses should implement this as:
	// - a normal method (so it's referentially shared on `prototype`)
	// - or assign a static function declaration.
	// Using an arrow func prop will cause recomputation on every render.
	// TODO: warn on successive recomputes!
	abstract produceStyles: SP

	// TODO: actually implement
	abstract get styleContext(): Context<MyTheme>

	get styles(): StaticSheet<ReturnType<SP>> {
		return globalStylesCache.getStaticSheet(
			this.produceStyles,
			this.styleContext,
			this.props
		)
	}
}

/**
 * UIButton
 * yikes.
 */
class UIButton<UIButtonProps> extends StyledComponent<UIButtonProps> {
	get stylesContext() {
		return globalContext
	}
	produceStyles(context: Context<MyTheme>) {
		return {
			button: {
				// Dang, these don't autocomplete....
			},
		}
	}
}

// Take 2: a more type-sympathetic approach
function createStyleProducer<Props, Classes extends string>(
	producer: (context: Context<MyTheme>) => Sheet<Props, Classes>
): (context: Context<MyTheme>) => Sheet<Props, Classes> {
	return producer
}

abstract class StyledComponent2<
	Props extends object,
	SP extends SheetProducer<MyTheme, Props, any>
> extends Component<Props> {
	abstract sheetProducer: SP

	get styles(): StaticSheet<ReturnType<SP>> {
		return globalStylesCache.getStaticSheet(
			this.sheetProducer,
			this.styleContext,
			this.props
		)
	}

	get styleContext(): Context<MyTheme> {
		return globalContext
	}
}

interface UIButtonProps {
	onClick?: (e: ReactMouseEvent<any, any>) => void
	isSecondary?: boolean
}

// Everything infers ok, but you need to specify the dynamic property function's argument type.
// This works great, even without `use const`.
const uiButtonStyles = createStyleProducer(({ theme }) => {
	return {
		button: {
			textAlign: "center",
			fontSize: 18,
			padding: "5px 10px",
			// Make single property dynamic.
			color:
				theme.name === "light"
					? ({ isSecondary }: UIButtonProps) => (isSecondary ? "#444" : "#000")
					: ({ isSecondary }: UIButtonProps) => (isSecondary ? "#ddd" : "#FFF"),
			// Override multiple properties at once
			_: ({ isSecondary }: UIButtonProps) => ({
				background: theme.name === "light" ? "#ddd" : "444",
				border: theme.name === "light" ? "2px solid black" : "2px solid white",
			}),
		},
	}
})

class UIButton2 extends StyledComponent2<UIButtonProps, typeof uiButtonStyles> {
	sheetProducer = uiButtonStyles
	render() {
		return (
			<div style={this.styles.button} onClick={this.props.onClick}>
				{this.props.children}
			</div>
		)
	}
}

// How can we make things like the StylesCache easier to build?
// It seems quite common to want a multi-level WeakMap-based cache
// for problems like memoization.

type WeakMapTreeData<T> = {
	map?: WeakMap<object, WeakMapTreeData<T>>
	value?: T
}

function createWeakMapTreeData<T>(): WeakMapTreeData<T> {
	return { map: new WeakMap() }
}

/**
 * WeakMapTree is a view into a (possibly-shared) WeakMapTreeData at a specific
 * level of heirarchy.
 *
 * The `Path` type paramter should be a fixed-length tuple type.
 * A path indexes into the WeakMapTreeData.
 *
 * For example, one could use a WeakMapTree to memoize a function with
 * three arguments:
 *
 * ```
 * function expensive(a: A, b: B, c: C): R
 * const expensiveCache = new WeakMapTree<[A, B, C], R>()
 * const expensiveMemoized = (a: A, b: B, c: C) => expensiveCache.compute([a, b, c], () => expensive(a, b, c))
 * ```
 */
class WeakMapTree<Path extends object[], Value> {
	public readonly data: WeakMapTreeData<Value>

	constructor(tree: WeakMapTreeData<Value> = createWeakMapTreeData()) {
		this.data = tree
	}

	set(path: Path, val: Value) {
		let data = this.data

		for (const part of path) {
			if (!data.map) {
				data.map = new WeakMap()
			}

			let nextData = data.map.get(part)
			if (!nextData) {
				nextData = { map: new WeakMap() }
				data.map.set(part, nextData)
			}

			data = nextData
		}

		data.value = val
	}

	get(path: Path): Value | undefined {
		let data: WeakMapTreeData<Value> | undefined = this.data
		for (const part of path) {
			if (data && data.map) {
				data = data.map.get(part)
			} else {
				return undefined
			}
		}

		return data && data.value
	}

	/**
	 * Retrieve the value at `path`, or compute it and store the
	 * value.
	 */
	compute(path: Path, valThunk: () => Value) {
		let value = this.get(path)
		if (!value) {
			value = valThunk()
			this.set(path, value)
		}
		return value
	}
}

// Now that we have WeakMapTree, we can use it to greatly simplify the implementation
// of StylesCache.
//
// Because of the WeakMapTree abstraction, it'd be easy to add or remove new layers
// of caching - for example, we could easily drop the props-related features.

class StylesCache2<Theme extends object> {
	tree = createWeakMapTreeData<any>()
	byProducer = new WeakMapTree<[SheetProducer<Theme, any, any>], unknown>(
		this.tree
	)
	sheetByProducerAndContext = new WeakMapTree<
		[SheetProducer<Theme, any, any>, Context<Theme>],
		Sheet<any, any>
	>(this.tree)
	staticSheetByProducerAndContextAndProps = new WeakMapTree<
		[SheetProducer<Theme, any, any>, Context<Theme>, object],
		StaticSheet<any>
	>(this.tree)

	getStaticSheet<Props extends object>(
		producer: SheetProducer<Theme, Props, any>,
		context: Context<Theme>,
		props: Props
	): StaticSheet<any> {
		const sheet = this.sheetByProducerAndContext.compute(
			[producer, context],
			() => producer(context)
		)

		return this.staticSheetByProducerAndContextAndProps.compute(
			[producer, context, props],
			() => getStaticSheetFromSheetAndProps(sheet, props)
		)
	}
}