Skip to content

Commit 60f92e6

Browse files
mgmeyersmgmeyers
committed
Improve section word count widget positioning
1 parent 74a8716 commit 60f92e6

File tree

2 files changed

+201
-165
lines changed

2 files changed

+201
-165
lines changed

‎src/editor/EditorPlugin.ts‎

Lines changed: 190 additions & 124 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { RangeSetBuilder, StateEffect, StateField, Transaction } from "@codemirror/state";
1+
import { Line, RangeSetBuilder, StateEffect, StateField, Text, Transaction } from "@codemirror/state";
22
import {
33
ViewUpdate,
44
PluginValue,
@@ -8,7 +8,6 @@ import {
88
Decoration,
99
WidgetType,
1010
} from "@codemirror/view";
11-
import { App, HeadingCache, TFile, editorInfoField } from "obsidian";
1211
import type BetterWordCount from "src/main";
1312
import { getWordCount } from "src/utils/StatUtils";
1413

@@ -78,161 +77,228 @@ class StatusBarEditorPlugin implements PluginValue {
7877

7978
export const statusBarEditorPlugin = ViewPlugin.fromClass(StatusBarEditorPlugin);
8079

81-
interface HeadingRange {
82-
heading: HeadingCache;
83-
from: number;
84-
to: number;
80+
interface SectionCountData {
81+
line: number;
82+
level: number;
83+
self: number;
84+
total: number;
85+
pos: number;
8586
}
8687

87-
function getHeadingRanges(app: App, file: TFile, end: number) {
88-
const fileCache = app.metadataCache.getFileCache(file);
89-
90-
if (!fileCache?.headings?.length) return null;
91-
92-
const nestedHeadings: HeadingCache[] = [];
93-
const ranges: HeadingRange[] = [];
94-
95-
for (let i = 0, len = fileCache.headings.length; i < len; i++) {
96-
const heading = fileCache.headings[i];
97-
const lastHeading = nestedHeadings.last();
98-
const isLast = i === len - 1;
99-
100-
if (!lastHeading || heading.level > lastHeading.level) {
101-
// First heading, or traversing to higher level heading (eg ## -> ###)
102-
nestedHeadings.push(heading);
103-
} else if (heading.level === lastHeading.level) {
104-
// Two headings of the same level
105-
const nestedHeading = nestedHeadings.pop();
106-
ranges.push({
107-
heading: nestedHeading,
108-
from: nestedHeading.position.end.offset,
109-
to: heading.position.start.offset,
110-
});
111-
nestedHeadings.push(heading);
112-
} else if (heading.level < lastHeading.level) {
113-
// Traversing to lower level heading (eg. ### -> ##)
114-
for (let j = nestedHeadings.length - 1; j >= 0; j--) {
115-
const nestedHeading = nestedHeadings[j];
116-
117-
if (heading.level < nestedHeading.level) {
118-
// Continue traversing to lower level heading
119-
const nestedHeading = nestedHeadings.pop();
120-
ranges.push({
121-
heading: nestedHeading,
122-
from: nestedHeading.position.end.offset,
123-
to: heading.position.start.offset,
124-
});
125-
if (j === 0) {
126-
nestedHeadings.push(heading);
127-
}
128-
continue;
129-
}
88+
class SectionWidget extends WidgetType {
89+
plugin: BetterWordCount;
90+
data: SectionCountData;
13091

131-
if (heading.level === nestedHeading.level) {
132-
// Stop because we found an equal level heading
133-
const nestedHeading = nestedHeadings.pop();
134-
ranges.push({
135-
heading: nestedHeading,
136-
from: nestedHeading.position.end.offset,
137-
to: heading.position.start.offset,
138-
});
139-
nestedHeadings.push(heading);
140-
break;
141-
}
92+
constructor(plugin: BetterWordCount, data: SectionCountData) {
93+
super();
94+
this.plugin = plugin;
95+
this.data = data;
96+
}
14297

143-
if (heading.level > nestedHeading.level) {
144-
// Stop because we found an higher level heading
145-
nestedHeadings.push(heading);
146-
break;
147-
}
148-
}
149-
} else if (isLast) {
150-
// Final heading
151-
nestedHeadings.push(heading);
152-
}
98+
eq(widget: this): boolean {
99+
const { pos, self, total } = this.data;
100+
return pos === widget.data.pos && self === widget.data.self && total === widget.data.total;
101+
}
153102

154-
if (isLast) {
155-
// Flush the remaining headings
156-
let nestedHeading: HeadingCache;
157-
while ((nestedHeading = nestedHeadings.pop())) {
158-
ranges.push({
159-
heading: nestedHeading,
160-
from: nestedHeading.position.end.offset,
161-
to: end,
162-
});
163-
}
103+
getDisplayText() {
104+
const { self, total } = this.data;
105+
if (self && self !== total) {
106+
return `${self} / ${total}`;
164107
}
108+
return total.toString();
165109
}
166110

167-
// Sort the headings in the order they appear in the document
168-
ranges.sort((a, b) => a.from - b.from);
169-
170-
return ranges;
111+
toDOM() {
112+
return createSpan({ cls: "bwc-section-count", text: this.getDisplayText() });
113+
}
171114
}
172115

173116
class SectionWordCountEditorPlugin implements PluginValue {
174117
decorations: DecorationSet;
118+
lineCounts: any[] = [];
175119

176120
constructor(view: EditorView) {
121+
const plugin = view.state.field(pluginField);
122+
if (!plugin.settings.displaySectionCounts) {
123+
this.decorations = Decoration.none;
124+
return;
125+
}
126+
127+
this.calculateLineCounts(view.state.doc);
177128
this.decorations = this.mkDeco(view);
178129
}
179130

131+
calculateLineCounts(doc: Text) {
132+
for (let index = 0; index < doc.lines; index++) {
133+
const line = doc.line(index + 1);
134+
this.lineCounts.push(getWordCount(line.text));
135+
}
136+
}
137+
180138
update(update: ViewUpdate) {
181139
const plugin = update.view.state.field(pluginField);
182-
if (!plugin.settings.displaySectionCounts) {
183-
if (this.decorations.size) {
184-
// Clear out any decorations
185-
this.decorations = this.mkDeco(update.view);
186-
}
140+
const { displaySectionCounts } = plugin.settings;
141+
let didSettingsChange = false;
187142

143+
if (this.lineCounts.length && !displaySectionCounts) {
144+
this.lineCounts = [];
145+
this.decorations = Decoration.none;
188146
return;
147+
} else if (!this.lineCounts.length && displaySectionCounts) {
148+
didSettingsChange = true;
149+
this.calculateLineCounts(update.startState.doc);
189150
}
190151

191-
update.transactions.forEach((tr) => {
192-
if (tr.effects.some((e) => e.is(metadataUpdated))) {
193-
// If metadata has been updated, rebuild the decorations
194-
this.decorations = this.mkDeco(update.view);
195-
} else if (update.docChanged) {
196-
// Otherwise just update their positions
197-
this.decorations = this.decorations.map(tr.changes);
198-
}
199-
});
152+
if (update.docChanged) {
153+
const startDoc = update.startState.doc;
154+
let tempDoc = startDoc;
155+
156+
update.changes.iterChanges((fromA, toA, fromB, toB, text) => {
157+
const from = fromB;
158+
const to = fromB + (toA - fromA);
159+
const nextTo = from + text.length;
160+
161+
const fromLine = tempDoc.lineAt(from);
162+
const toLine = tempDoc.lineAt(to);
163+
164+
tempDoc = tempDoc.replace(fromB, fromB + (toA - fromA), text);
165+
166+
const fromLineNext = tempDoc.lineAt(from);
167+
const toLineNext = tempDoc.lineAt(nextTo);
168+
169+
const lines: any[] = [];
170+
171+
for (let i = fromLineNext.number; i <= toLineNext.number; i++) {
172+
lines.push(getWordCount(tempDoc.line(i).text));
173+
}
174+
175+
const spliceStart = fromLine.number - 1;
176+
const spliceLen = toLine.number - fromLine.number + 1;
177+
178+
this.lineCounts.splice(spliceStart, spliceLen, ...lines);
179+
});
180+
}
181+
182+
if (update.docChanged || update.viewportChanged || didSettingsChange) {
183+
this.decorations = this.mkDeco(update.view);
184+
}
200185
}
201186

202187
mkDeco(view: EditorView) {
203188
const plugin = view.state.field(pluginField);
204189
const b = new RangeSetBuilder<Decoration>();
205190
if (!plugin.settings.displaySectionCounts) return b.finish();
206191

207-
const { app, file } = view.state.field(editorInfoField);
208-
if (!file) return b.finish();
209-
210-
const headingRanges = getHeadingRanges(app, file, view.state.doc.length - 1);
211-
if (!headingRanges?.length) return b.finish();
212-
213-
for (let i = 0; i < headingRanges.length; i++) {
214-
const heading = headingRanges[i];
215-
const next = headingRanges[i + 1];
216-
const targetPos = heading.heading.position.start.offset;
192+
const getHeaderLevel = (line: Line) => {
193+
const match = line.text.match(/^(#+)[ \t]/);
194+
return match ? match[1].length : null;
195+
};
196+
197+
const doc = view.state.doc;
198+
const lineCount = doc.lines;
199+
const sectionCounts: SectionCountData[] = [];
200+
const nested: SectionCountData[] = [];
201+
202+
for (const { from } of view.visibleRanges) {
203+
const lineStart = doc.lineAt(from);
204+
205+
for (let i = lineStart.number, len = lineCount; i <= len; i++) {
206+
let line: Line;
207+
if (i === lineStart.number) line = lineStart;
208+
else line = doc.line(i);
209+
210+
const level = getHeaderLevel(line);
211+
const prevHeading = nested.last();
212+
if (level) {
213+
if (!prevHeading || level > prevHeading.level) {
214+
nested.push({
215+
line: i,
216+
level,
217+
self: 0,
218+
total: 0,
219+
pos: line.to,
220+
});
221+
} else if (prevHeading.level === level) {
222+
const nestedHeading = nested.pop();
223+
sectionCounts.push(nestedHeading);
224+
nested.push({
225+
line: i,
226+
level,
227+
self: 0,
228+
total: 0,
229+
pos: line.to,
230+
});
231+
} else if (prevHeading.level > level) {
232+
// Traversing to lower level heading (eg. ### -> ##)
233+
for (let j = nested.length - 1; j >= 0; j--) {
234+
const nestedHeading = nested[j];
235+
236+
if (level < nestedHeading.level) {
237+
// Continue traversing to lower level heading
238+
const nestedHeading = nested.pop();
239+
sectionCounts.push(nestedHeading);
240+
if (j === 0) {
241+
nested.push({
242+
line: i,
243+
level,
244+
self: 0,
245+
total: 0,
246+
pos: line.to,
247+
});
248+
}
249+
continue;
250+
}
251+
252+
if (level === nestedHeading.level) {
253+
// Stop because we found an equal level heading
254+
const nestedHeading = nested.pop();
255+
sectionCounts.push(nestedHeading);
256+
nested.push({
257+
line: i,
258+
level,
259+
self: 0,
260+
total: 0,
261+
pos: line.to,
262+
});
263+
break;
264+
}
265+
266+
if (level > nestedHeading.level) {
267+
// Stop because we found an higher level heading
268+
nested.push({
269+
line: i,
270+
level,
271+
self: 0,
272+
total: 0,
273+
pos: line.to,
274+
});
275+
break;
276+
}
277+
}
278+
}
279+
} else if (nested.length) {
280+
const count = this.lineCounts[i - 1];
281+
for (const heading of nested) {
282+
if (heading === prevHeading) {
283+
heading.self += count;
284+
}
285+
heading.total += count;
286+
}
287+
}
288+
}
289+
}
217290

218-
const totalCount = getWordCount(view.state.doc.slice(heading.from, heading.to).toString());
219-
let selfCount: number;
291+
if (nested.length) sectionCounts.push(...nested);
220292

221-
if (next && next.heading.level > heading.heading.level) {
222-
const betweenCount = getWordCount(
223-
view.state.doc.slice(heading.from, next.heading.position.start.offset).toString()
224-
);
225-
if (betweenCount) selfCount = betweenCount;
226-
}
293+
sectionCounts.sort((a, b) => a.line - b.line);
227294

295+
for (const data of sectionCounts) {
228296
b.add(
229-
targetPos,
230-
targetPos,
231-
Decoration.line({
232-
attributes: {
233-
class: "bwc-section-count",
234-
style: `--word-count: "${selfCount ? `${selfCount} / ${totalCount}` : totalCount.toString()}"`,
235-
},
297+
data.pos,
298+
data.pos,
299+
Decoration.widget({
300+
side: 1,
301+
widget: new SectionWidget(plugin, data),
236302
})
237303
);
238304
}
@@ -241,7 +307,7 @@ class SectionWordCountEditorPlugin implements PluginValue {
241307
}
242308
}
243309

244-
export const metadataUpdated = StateEffect.define<void>();
310+
export const settingsChanged = StateEffect.define<void>();
245311
export const sectionWordCountEditorPlugin = ViewPlugin.fromClass(SectionWordCountEditorPlugin, {
246312
decorations: (v) => v.decorations,
247313
});

0 commit comments

Comments
 (0)