Skip to content

Commit 3e5b3ee

Browse files
fredericbahrfredericbahr
committed
refactor: implement subscribe as a lit directive
An async lit directive allows for fine grained reactivity as it only updates when an external source (here the atom/store) is updated. This aligns closely with the fine-grained reactivity model we have in react.
1 parent bb629bf commit 3e5b3ee

3 files changed

Lines changed: 202 additions & 136 deletions

File tree

‎packages/lit-table/src/TableController.ts‎

Lines changed: 27 additions & 136 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,16 @@
11
import { constructTable } from '@tanstack/table-core'
2-
import { TanStackStoreSelector } from '@tanstack/lit-store'
32
import { litReactivity } from './reactivity'
43
import { FlexRender } from './flexRender'
4+
5+
import { subscribe } from './subscribe-directive'
56
import type {
6-
Atom,
7-
ReadonlyAtom,
8-
ReadonlyStore,
9-
Store,
10-
} from '@tanstack/lit-store'
11-
import type {
12-
NoInfer,
137
RowData,
148
Table,
159
TableFeatures,
1610
TableOptions,
1711
TableState,
1812
} from '@tanstack/table-core'
19-
import type {
20-
ReactiveController,
21-
ReactiveControllerHost,
22-
TemplateResult,
23-
} from 'lit'
24-
25-
export type SubscribeSource<TValue> =
26-
| Atom<TValue>
27-
| ReadonlyAtom<TValue>
28-
| Store<TValue>
29-
| ReadonlyStore<TValue>
13+
import type { ReactiveController, ReactiveControllerHost } from 'lit'
3014

3115
/**
3216
* The extended table type returned by the Lit adapter.
@@ -46,45 +30,34 @@ export type LitTable<
4630
*/
4731
readonly store: Table<TFeatures, TData>['store']
4832
/**
49-
* Subscribe to a selected slice of table state, or to a single source (atom or store).
50-
*
51-
* **Lit note:** `TableController` still wires host updates via the full `table.store`
52-
* subscription — source mode matches the React API and reads `source.get()` at render
53-
* time. True source-only invalidation can be added later via `source.subscribe`.
33+
* Subscribes to the table's underlying state store within a Lit template.
34+
* Re-renders only the targeted template slice when the observed state changes.
5435
*
5536
* @example
5637
* ```ts
57-
* table.Subscribe({
58-
* selector: (state) => ({ rowSelection: state.rowSelection }),
59-
* children: (state) => html`<div>${JSON.stringify(state)}</div>`,
60-
* })
38+
* // 1. Subscribe to a specific state slice (re-renders ONLY when rowSelection changes)
39+
* html`
40+
* <div>
41+
* ${table.subscribe(
42+
* table.store,
43+
* (state) => state.rowSelection,
44+
* (rowSelection) => html`<span>Selected: ${JSON.stringify(rowSelection)}</span>`
45+
* )}
46+
* </div>
47+
* `
48+
*
49+
* // 2. Subscribe to the full state (re-renders on any state mutation)
50+
* html`
51+
* <div>
52+
* ${table.subscribe(
53+
* table.store,
54+
* (state) => html`<span>Total rows: ${state.rowModel.rows.length}</span>`
55+
* )}
56+
* </div>
57+
* `
6158
* ```
6259
*/
63-
Subscribe: {
64-
<TSourceValue>(props: {
65-
source: SubscribeSource<TSourceValue>
66-
selector?: undefined
67-
children:
68-
| ((state: Readonly<TSourceValue>) => TemplateResult | string)
69-
| TemplateResult
70-
| string
71-
}): TemplateResult | string
72-
<TSourceValue, TSubscribeSelected>(props: {
73-
source: SubscribeSource<TSourceValue>
74-
selector: (state: TSourceValue) => TSubscribeSelected
75-
children:
76-
| ((state: Readonly<TSubscribeSelected>) => TemplateResult | string)
77-
| TemplateResult
78-
| string
79-
}): TemplateResult | string
80-
<TSubscribeSelected>(props: {
81-
selector: (state: NoInfer<TableState<TFeatures>>) => TSubscribeSelected
82-
children:
83-
| ((state: Readonly<TSubscribeSelected>) => TemplateResult | string)
84-
| TemplateResult
85-
| string
86-
}): TemplateResult | string
87-
}
60+
subscribe: typeof subscribe
8861
/**
8962
* The selected state of the table. This state may not match the structure of
9063
* the full table state because it is selected by the selector function that
@@ -152,13 +125,6 @@ export class TableController<
152125
private _storeSubscription?: { unsubscribe: () => void }
153126
private _optionsSubscription?: { unsubscribe: () => void }
154127
private _notifier = 0
155-
private _selectorCache = new WeakMap<
156-
SubscribeSource<unknown>,
157-
Map<
158-
((state: unknown) => unknown) | undefined,
159-
TanStackStoreSelector<unknown>
160-
>
161-
>()
162128

163129
constructor(host: ReactiveControllerHost) {
164130
;(this.host = host).addController(this)
@@ -217,33 +183,9 @@ export class TableController<
217183
// Capture for closure
218184
const tableInstance = this._table
219185

220-
// Attach Subscribe function
221-
const Subscribe = ((props: {
222-
source?: SubscribeSource<unknown>
223-
selector?: (state: unknown) => unknown
224-
children:
225-
| ((state: Readonly<unknown>) => TemplateResult | string)
226-
| TemplateResult
227-
| string
228-
}): TemplateResult | string => {
229-
const source = props.source ?? tableInstance.store
230-
231-
const storeSelector: TanStackStoreSelector<unknown> =
232-
this._getOrCreateSelector(source, props.selector)
233-
234-
// TODO: update to newest version of Tanstack Store: https://github.com/TanStack/store/pull/329
235-
const selectedState = storeSelector.value
236-
237-
if (typeof props.children === 'function') {
238-
return props.children(selectedState as Readonly<unknown>)
239-
}
240-
241-
return props.children
242-
}) as LitTable<TFeatures, TData, TSelected>['Subscribe']
243-
244186
return {
245187
...this._table,
246-
Subscribe,
188+
subscribe,
247189
FlexRender,
248190
get state() {
249191
return (selector?.(tableInstance.store.state) ??
@@ -275,56 +217,5 @@ export class TableController<
275217
this._storeSubscription = undefined
276218
this._optionsSubscription?.unsubscribe()
277219
this._optionsSubscription = undefined
278-
this._selectorCache = new WeakMap()
279-
}
280-
281-
/**
282-
* Get or create a TanStackStoreSelector for the given source and selector.
283-
*
284-
* Caches selectors by source (WeakMap) and selector function to avoid
285-
* creating new controllers on every render cycle.
286-
*
287-
* @param source The atom or store to subscribe to
288-
* @param selector Optional selector function to select a slice of the source state
289-
* @returns A cached TanStackStoreSelector instance that subscribes to the source and applies the selector
290-
*/
291-
private _getOrCreateSelector = (
292-
source?: SubscribeSource<unknown>,
293-
selector?: (state: unknown) => unknown,
294-
): TanStackStoreSelector<unknown> => {
295-
if (!source) {
296-
return new TanStackStoreSelector(this.host, () => source, selector)
297-
}
298-
299-
if (!this._selectorCache.has(source)) {
300-
this._selectorCache.set(source, new Map())
301-
}
302-
const selectorMap = this._selectorCache.get(source)
303-
304-
// Get or create the selector for this source + selector combination
305-
if (selectorMap?.has(selector)) {
306-
return (
307-
selectorMap.get(selector) ??
308-
this.createSelectorForSource(source, selector)
309-
)
310-
}
311-
312-
const storeSelector = this.createSelectorForSource(source, selector)
313-
selectorMap?.set(selector, storeSelector)
314-
315-
return storeSelector
316-
}
317-
318-
/**
319-
* Create a new TanStackStoreSelector for the given source and selector without caching.
320-
* @param source The atom or store to subscribe to
321-
* @param selector Optional selector function to select a slice of the source state
322-
* @returns A new TanStackStoreSelector instance that subscribes to the source and applies the selector
323-
*/
324-
private createSelectorForSource = (
325-
source: SubscribeSource<unknown>,
326-
selector?: (state: unknown) => unknown,
327-
): TanStackStoreSelector<unknown> => {
328-
return new TanStackStoreSelector(this.host, () => source, selector)
329220
}
330221
}

‎packages/lit-table/src/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,5 @@ export * from '@tanstack/table-core'
22

33
export * from './flexRender'
44
export * from './TableController'
5+
export * from './subscribe-directive'
56
export * from './createTableHook'

‎packages/lit-table/src/subscribe-directive.ts‎

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
import { TanStackStoreSelector } from '@tanstack/lit-store'
2+
import { AsyncDirective, directive } from 'lit/async-directive.js'
3+
import type { DirectiveResult } from 'lit/async-directive.js'
4+
import type {
5+
Atom,
6+
ReadonlyAtom,
7+
ReadonlyStore,
8+
Store,
9+
} from '@tanstack/lit-store'
10+
import type { ReactiveControllerHost } from 'lit'
11+
12+
export type SelectionSource<TValue> =
13+
| Atom<TValue>
14+
| ReadonlyAtom<TValue>
15+
| Store<TValue>
16+
| ReadonlyStore<TValue>
17+
18+
/**
19+
* A function that selects a specific slice of state from the source.
20+
* @template TSource - The complete state type from the store/atom.
21+
* @template TSelected - The extracted or derived state type.
22+
*/
23+
type Selector<TSource, TSelected> = (state: TSource) => TSelected
24+
25+
/**
26+
* A render function that takes the selected state and returns content
27+
* (typically a `TemplateResult`) to be rendered by Lit.
28+
* @template TSelected - The selected state passed into the template.
29+
*/
30+
type TemplateFunction<TSelected> = (value: TSelected) => unknown
31+
32+
/**
33+
* A simple identity selector used when no specific selection is needed,
34+
* allowing the directive to subscribe to the entire state of the source.
35+
* @template T - The type of the state being passed through unchanged.
36+
*/
37+
const identitySelector = <T>(state: T): T => state
38+
39+
/**
40+
* An asynchronous Lit directive that subscribes to a `@tanstack/lit-store`
41+
* source and triggers re-renders specifically for the template portion it wraps.
42+
* * It uses a "fake" `ReactiveControllerHost` to bridge the gap between
43+
* TanStack's standard controller requirements and the `AsyncDirective` lifecycle.
44+
*/
45+
export class SubscribeDirective extends AsyncDirective {
46+
/** The `TanStackStoreSelector` controller that manages the subscription to the store/atom */
47+
private controller?: TanStackStoreSelector<any, any>
48+
49+
/** The latest source and selector used to determine if a new subscription is needed on updates */
50+
private latestSource?: SelectionSource<any>
51+
/* The latest selector function used to determine if a new subscription is needed on updates */
52+
private latestSelector?: Selector<any, any>
53+
/* The latest resolved template function to render the selected state slice */
54+
private resolvedTemplate?: TemplateFunction<any>
55+
56+
/**
57+
* Renders the entire state of the source without a selector.
58+
* @param source - The store or atom to subscribe to.
59+
* @param template - The render function receiving the full state.
60+
*/
61+
render<TSource>(
62+
source: SelectionSource<TSource>,
63+
template: TemplateFunction<TSource>,
64+
): unknown
65+
66+
/**
67+
* Renders a specific slice of state derived via a selector function.
68+
* @param source - The store or atom to subscribe to.
69+
* @param selector - A function to extract the relevant slice of state.
70+
* @param template - The render function receiving the selected state slice.
71+
*/
72+
render<TSource, TSelected>(
73+
source: SelectionSource<TSource>,
74+
selector: Selector<TSource, TSelected>,
75+
template: TemplateFunction<TSelected>,
76+
): unknown
77+
78+
render(
79+
source: SelectionSource<any>,
80+
selectorOrTemplate: Selector<any, any> | TemplateFunction<any>,
81+
template?: TemplateFunction<any>,
82+
) {
83+
const isIdentitySubscription: boolean = template === undefined
84+
85+
const selector = isIdentitySubscription
86+
? identitySelector
87+
: (selectorOrTemplate as Selector<any, any>)
88+
89+
const actualTemplate = isIdentitySubscription
90+
? (selectorOrTemplate as TemplateFunction<any>)
91+
: template
92+
93+
if (this.latestSelector !== selector) {
94+
this.controller?.hostDisconnected()
95+
this.controller = undefined
96+
}
97+
98+
this.latestSource = source
99+
this.latestSelector = selector
100+
this.resolvedTemplate = actualTemplate
101+
102+
if (!this.controller) {
103+
this.controller = new TanStackStoreSelector(
104+
this.createFakeHost(),
105+
() => this.latestSource,
106+
(state) => this.latestSelector?.(state),
107+
)
108+
}
109+
110+
this.controller.hostUpdate()
111+
112+
// TODO: update to newest version of Tanstack Store: https://github.com/TanStack/store/pull/329
113+
return this.resolvedTemplate?.(this.controller.value)
114+
}
115+
116+
/** Cleans up the controller subscription when the directive is removed from the DOM. */
117+
disconnected() {
118+
this.controller?.hostDisconnected()
119+
}
120+
121+
/** Restores the controller subscription when the directive is re-attached to the DOM. */
122+
reconnected() {
123+
this.controller?.hostUpdate()
124+
}
125+
126+
/**
127+
* Creates a mock `ReactiveControllerHost` allowing the `TanStackStoreSelector`
128+
* to plug into the `AsyncDirective`'s update cycle using `setValue()`.
129+
*/
130+
private createFakeHost(): ReactiveControllerHost {
131+
return {
132+
addController: () => {},
133+
removeController: () => {},
134+
requestUpdate: () => {
135+
if (this.resolvedTemplate && this.controller) {
136+
this.setValue(this.resolvedTemplate(this.controller.value))
137+
}
138+
},
139+
get updateComplete() {
140+
return Promise.resolve(true)
141+
},
142+
}
143+
}
144+
}
145+
146+
/**
147+
* A Lit directive that subscribes to a source (Store or Atom)
148+
* and efficiently updates only the wrapped template
149+
* when the state or selected slice changes.
150+
* @example
151+
* ```ts
152+
* // Without a selector (subscribes to entire state)
153+
* html`<div>${subscribe(myStore, (state) => html`<span>${state.count}</span>`)}</div>`
154+
* * // With a selector (only updates when `count` changes)
155+
* html`<div>${subscribe(myStore, state => state.count, (count) => html`<span>${count}</span>`)}</div>`
156+
* ```
157+
*/
158+
export const subscribe = directive(SubscribeDirective) as {
159+
/** Subscribes to the entire source state without filtering. */
160+
<TSource>(
161+
source: SelectionSource<TSource>,
162+
template: TemplateFunction<TSource>,
163+
): DirectiveResult<typeof SubscribeDirective>
164+
165+
/**
166+
* Subscribes to a specific slice of the source state via a selector,
167+
* preventing unnecessary re-renders when other parts of the state change.
168+
*/
169+
<TSource, TSelected>(
170+
source: SelectionSource<TSource>,
171+
selector: Selector<TSource, TSelected>,
172+
template: TemplateFunction<TSelected>,
173+
): DirectiveResult<typeof SubscribeDirective>
174+
}

0 commit comments

Comments
 (0)