Repository navigation
Expand file tree
/
Copy pathrender_state.lua
More file actions
683 lines (597 loc) · 18.3 KB
/
Copy pathrender_state.lua
File metadata and controls
683 lines (597 loc) · 18.3 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
---@class RenderedMessage
---@field message OpencodeMessage Direct reference to message in state.messages
---@field line_start integer? Line where message header starts
---@field line_end integer? Line where message header ends
---@class RenderedPart
---@field part OpencodeMessagePart Direct reference to part in state.messages
---@field message_id string ID of parent message
---@field line_start integer? Line where part starts
---@field line_end integer? Line where part ends
---@field actions table[] Actions associated with this part
---@field has_extmarks boolean? Whether the part currently has extmarks applied
---@class RenderState
---@field _messages table<string, RenderedMessage> Message ID -> rendered message
---@field _parts table<string, RenderedPart> Part ID -> rendered part
---@field _part_ranges {[1]: integer, [2]: integer, [3]: string}[] Sorted [line_start, line_end, part_id] for binary search
---@field _message_ranges {[1]: integer, [2]: integer, [3]: string}[] Sorted [line_start, line_end, message_id] for binary search
---@field _ranges_valid boolean Whether range arrays are sorted and up-to-date
---@field _max_line_end integer
---@field _max_line_end_valid boolean
local RenderState = {}
RenderState.__index = RenderState
---@return RenderState
function RenderState.new()
local self = setmetatable({}, RenderState)
self:reset()
return self
end
function RenderState:reset()
self._messages = {}
self._parts = {}
self._orphan_parts = {}
self._orphan_parts_index = {}
self._part_ranges = {}
self._message_ranges = {}
self._ranges_valid = false
self._max_line_end = 0
self._max_line_end_valid = true
self._child_session_parts = {}
self._child_session_parts_index = {} -- session_id -> part_id -> list_index
self._child_session_task_parts = {}
self._task_part_child_sessions = {}
self._snapshot_id_index = {} -- snapshot_id -> OpencodeMessagePart
end
function RenderState:_recompute_max_line_end()
local max_line_end = 0
for _, msg_data in pairs(self._messages) do
if msg_data.line_end and msg_data.line_end > max_line_end then
max_line_end = msg_data.line_end
end
end
for _, part_data in pairs(self._parts) do
if part_data.line_end and part_data.line_end > max_line_end then
max_line_end = part_data.line_end
end
end
self._max_line_end = max_line_end
self._max_line_end_valid = true
return max_line_end
end
---@return integer
function RenderState:_get_max_line_end()
if not self._max_line_end_valid then
return self:_recompute_max_line_end()
end
return self._max_line_end
end
---@param part OpencodeMessagePart?
---@return string?
local function get_child_session_id_for_task_part(part)
if not part or part.tool ~= 'task' then
return nil
end
local part_state = part.state
local metadata = part_state and part_state.metadata
return metadata and metadata.sessionId or nil
end
---@param part_id string
function RenderState:_clear_task_part_child_session(part_id)
local child_session_id = self._task_part_child_sessions[part_id]
if not child_session_id then
return
end
if self._child_session_task_parts[child_session_id] == part_id then
self._child_session_task_parts[child_session_id] = nil
end
self._task_part_child_sessions[part_id] = nil
end
---@param part_id string
---@param part OpencodeMessagePart
function RenderState:_index_task_part_child_session(part_id, part)
self:_clear_task_part_child_session(part_id)
local child_session_id = get_child_session_id_for_task_part(part)
if not child_session_id then
return
end
self._child_session_task_parts[child_session_id] = part_id
self._task_part_child_sessions[part_id] = child_session_id
end
---@param ranges {[1]: integer, [2]: integer, [3]: string}[]
---@param line integer
---@return string?
local function range_lookup(ranges, line)
local lo, hi = 1, #ranges
while lo <= hi do
local mid = math.floor((lo + hi) / 2)
local r = ranges[mid]
if line < r[1] then
hi = mid - 1
elseif line > r[2] then
lo = mid + 1
else
return r[3]
end
end
return nil
end
function RenderState:_rebuild_ranges()
local part_ranges = {}
for part_id, part_data in pairs(self._parts) do
if part_data.line_start and part_data.line_end then
part_ranges[#part_ranges + 1] = { part_data.line_start, part_data.line_end, part_id }
end
end
table.sort(part_ranges, function(a, b)
return a[1] < b[1]
end)
self._part_ranges = part_ranges
local msg_ranges = {}
for msg_id, msg_data in pairs(self._messages) do
if msg_data.line_start and msg_data.line_end then
msg_ranges[#msg_ranges + 1] = { msg_data.line_start, msg_data.line_end, msg_id }
end
end
table.sort(msg_ranges, function(a, b)
return a[1] < b[1]
end)
self._message_ranges = msg_ranges
self._ranges_valid = true
end
function RenderState:_ensure_ranges()
if not self._ranges_valid then
self:_rebuild_ranges()
end
end
---@param session_id string
---@return OpencodeMessagePart[]?
function RenderState:get_child_session_parts(session_id)
if not session_id then
return nil
end
return self._child_session_parts[session_id]
end
---@param session_id string
---@return string?
function RenderState:get_task_part_by_child_session(session_id)
if not session_id then
return nil
end
return self._child_session_task_parts[session_id]
end
---@param session_id string
---@param part OpencodeMessagePart
function RenderState:upsert_child_session_part(session_id, part)
if not session_id or not part or not part.id then
return
end
local session_parts = self._child_session_parts[session_id]
if not session_parts then
session_parts = {}
self._child_session_parts[session_id] = session_parts
self._child_session_parts_index[session_id] = {}
end
local idx = self._child_session_parts_index[session_id][part.id]
if idx then
session_parts[idx] = part
else
session_parts[#session_parts + 1] = part
self._child_session_parts_index[session_id][part.id] = #session_parts
end
end
---@param message_id string
---@return RenderedMessage?
function RenderState:get_message(message_id)
return self._messages[message_id]
end
---@param messages OpencodeMessage[]
---@param message_id string
---@return RenderedMessage?
function RenderState:get_previous_message(messages, message_id)
for i = #messages, 1, -1 do
local message = messages[i]
if message and message.info and message.info.id == message_id then
if i <= 1 then
return nil
end
local previous_message = messages[i - 1]
return previous_message and previous_message.info and self._messages[previous_message.info.id] or nil
end
end
return nil
end
---@param message_id string
---@param part OpencodeMessagePart
function RenderState:upsert_orphan_part(message_id, part)
if not message_id or not part or not part.id then
return
end
local orphan_parts = self._orphan_parts[message_id]
if not orphan_parts then
orphan_parts = {}
self._orphan_parts[message_id] = orphan_parts
self._orphan_parts_index[message_id] = {}
end
local orphan_index = self._orphan_parts_index[message_id]
local idx = orphan_index[part.id]
if idx then
orphan_parts[idx] = part
else
orphan_parts[#orphan_parts + 1] = part
orphan_index[part.id] = #orphan_parts
end
end
---@param message_id string
---@return OpencodeMessagePart[]
function RenderState:consume_orphan_parts(message_id)
if not message_id then
return {}
end
local orphan_parts = self._orphan_parts[message_id] or {}
self._orphan_parts[message_id] = nil
self._orphan_parts_index[message_id] = nil
return orphan_parts
end
---@param message_id string
---@param part_id string
---@return boolean
function RenderState:remove_orphan_part(message_id, part_id)
local orphan_parts = message_id and self._orphan_parts[message_id]
local orphan_index = message_id and self._orphan_parts_index[message_id]
local idx = orphan_index and orphan_index[part_id]
if not idx then
return false
end
table.remove(orphan_parts, idx)
orphan_index[part_id] = nil
for i = idx, #orphan_parts do
local part = orphan_parts[i]
if part and part.id then
orphan_index[part.id] = i
end
end
if #orphan_parts == 0 then
self._orphan_parts[message_id] = nil
self._orphan_parts_index[message_id] = nil
end
return true
end
---@param message_id string
function RenderState:clear_orphan_parts(message_id)
if not message_id then
return
end
self._orphan_parts[message_id] = nil
self._orphan_parts_index[message_id] = nil
end
---@param line integer 1-indexed
---@return RenderedMessage?
function RenderState:get_message_at_line(line)
self:_ensure_ranges()
local msg_id = range_lookup(self._message_ranges, line)
return msg_id and self._messages[msg_id] or nil
end
---@param part_id string
---@return RenderedPart?
function RenderState:get_part(part_id)
return self._parts[part_id]
end
---@param line integer 1-indexed
---@return RenderedPart?
function RenderState:get_part_at_line(line)
self:_ensure_ranges()
local part_id = range_lookup(self._part_ranges, line)
return part_id and self._parts[part_id] or nil
end
---@param call_id string
---@param message_id string
---@return string?
function RenderState:get_part_by_call_id(call_id, message_id)
local rendered_message = self._messages[message_id]
if rendered_message and rendered_message.message and rendered_message.message.parts then
for _, part in ipairs(rendered_message.message.parts) do
if part.callID == call_id then
return part.id
end
end
end
return nil
end
---@param snapshot_id string
---@return OpencodeMessagePart?
function RenderState:get_part_by_snapshot_id(snapshot_id)
return self._snapshot_id_index[snapshot_id]
end
---@param line integer
---@return table[]
function RenderState:get_actions_at_line(line)
self:_ensure_ranges()
local part_id = range_lookup(self._part_ranges, line)
if not part_id then
return {}
end
local part_data = self._parts[part_id]
if not part_data or not part_data.actions then
return {}
end
local actions = {}
for _, action in ipairs(part_data.actions) do
if action.range and action.range.from <= line and action.range.to >= line then
actions[#actions + 1] = action
end
end
return actions
end
---@param part_id string
---@param actions table[]
---@param offset? integer Line offset to apply to action line numbers
function RenderState:add_actions(part_id, actions, offset)
local part_data = self._parts[part_id]
if not part_data then
return
end
offset = offset or 0
for _, action in ipairs(actions) do
if offset ~= 0 then
if action.display_line then
action.display_line = action.display_line + offset
end
if action.range then
action.range.from = action.range.from + offset
action.range.to = action.range.to + offset
end
end
part_data.actions[#part_data.actions + 1] = action
end
end
---@param part_id string
function RenderState:clear_actions(part_id)
local part_data = self._parts[part_id]
if part_data then
part_data.actions = {}
end
end
---@return table[]
function RenderState:get_all_actions()
local all_actions = {}
for _, part_data in pairs(self._parts) do
if part_data.actions then
for _, action in ipairs(part_data.actions) do
all_actions[#all_actions + 1] = action
end
end
end
return all_actions
end
---@param message OpencodeMessage
---@param line_start integer?
---@param line_end integer?
function RenderState:set_message(message, line_start, line_end)
if not message or not message.info or not message.info.id then
return
end
local message_id = message.info.id
local existing = self._messages[message_id]
if not existing then
self._messages[message_id] = {
message = message,
line_start = line_start,
line_end = line_end,
}
else
existing.message = message
if line_start then
existing.line_start = line_start
end
if line_end then
existing.line_end = line_end
end
end
if line_start and line_end then
self._ranges_valid = false
if self._max_line_end_valid and line_end > self._max_line_end then
self._max_line_end = line_end
end
end
end
---@param part OpencodeMessagePart
---@param line_start integer?
---@param line_end integer?
function RenderState:set_part(part, line_start, line_end)
if not part or not part.id then
return
end
local part_id = part.id
local message_id = part.messageID or 'special'
local existing = self._parts[part_id]
if not existing then
self._parts[part_id] = {
part = part,
message_id = message_id,
line_start = line_start,
line_end = line_end,
actions = {},
has_extmarks = false,
}
else
existing.part = part
if message_id then
existing.message_id = message_id
end
if line_start then
existing.line_start = line_start
end
if line_end then
existing.line_end = line_end
end
end
if line_start and line_end then
self._ranges_valid = false
if self._max_line_end_valid and line_end > self._max_line_end then
self._max_line_end = line_end
end
end
if part.type == 'patch' and part.hash then
self._snapshot_id_index[part.hash] = part
end
self:_index_task_part_child_session(part_id, part)
end
---@param part_id string
---@param new_line_start integer
---@param new_line_end integer
---@return boolean
function RenderState:update_part_lines(part_id, new_line_start, new_line_end)
local part_data = self._parts[part_id]
if not part_data or not part_data.line_start or not part_data.line_end then
return false
end
if part_data.line_start == new_line_start and part_data.line_end == new_line_end then
return true
end
local old_line_end = part_data.line_end
local old_line_count = old_line_end - part_data.line_start + 1
local new_line_count = new_line_end - new_line_start + 1
local delta = new_line_count - old_line_count
part_data.line_start = new_line_start
part_data.line_end = new_line_end
self._ranges_valid = false
if self._max_line_end_valid then
if old_line_end == self._max_line_end and new_line_end < old_line_end then
self._max_line_end_valid = false
elseif new_line_end > self._max_line_end then
self._max_line_end = new_line_end
end
end
if delta ~= 0 then
self:shift_all(old_line_end + 1, delta)
end
return true
end
---@param part_ref OpencodeMessagePart
---@return RenderedPart?
function RenderState:update_part_data(part_ref)
if not part_ref or not part_ref.id then
return
end
local rendered_part = self._parts[part_ref.id]
if not rendered_part then
return
end
rendered_part.part = part_ref
if part_ref.type == 'patch' and part_ref.hash then
self._snapshot_id_index[part_ref.hash] = part_ref
end
self:_index_task_part_child_session(part_ref.id, part_ref)
return rendered_part
end
---@param part_id string
---@return boolean
function RenderState:remove_part(part_id)
local part_data = self._parts[part_id]
if not part_data then
return false
end
if part_data.part and part_data.part.type == 'patch' and part_data.part.hash then
self._snapshot_id_index[part_data.part.hash] = nil
end
self:_clear_task_part_child_session(part_id)
if not part_data.line_start or not part_data.line_end then
self._parts[part_id] = nil
return true
end
local line_count = part_data.line_end - part_data.line_start + 1
local shift_from = part_data.line_end + 1
self._parts[part_id] = nil
self._ranges_valid = false
if self._max_line_end_valid and part_data.line_end == self._max_line_end then
self._max_line_end_valid = false
end
self:shift_all(shift_from, -line_count)
return true
end
---@param message_id string
---@return boolean
function RenderState:remove_message(message_id)
local msg_data = self._messages[message_id]
if not msg_data or not msg_data.line_start or not msg_data.line_end then
return false
end
local line_count = msg_data.line_end - msg_data.line_start + 1
local shift_from = msg_data.line_end + 1
self._messages[message_id] = nil
self._ranges_valid = false
if self._max_line_end_valid and msg_data.line_end == self._max_line_end then
self._max_line_end_valid = false
end
self:shift_all(shift_from, -line_count)
return true
end
local function shift_action(action, delta)
if action.display_line then
action.display_line = action.display_line + delta
end
if action.range then
action.range.from = action.range.from + delta
action.range.to = action.range.to + delta
end
end
--- Binary-search sorted ranges for the first entry whose line_start >= from_line.
--- Returns the index, or #ranges + 1 if none qualify.
---@param ranges {[1]: integer, [2]: integer, [3]: string}[]
---@param from_line integer
---@return integer
local function first_range_at_or_after(ranges, from_line)
local lo, hi = 1, #ranges
local result = #ranges + 1
while lo <= hi do
local mid = math.floor((lo + hi) / 2)
if ranges[mid][1] >= from_line then
result = mid
hi = mid - 1
else
lo = mid + 1
end
end
return result
end
function RenderState:shift_all(from_line, delta)
if delta == 0 then
return
end
if from_line > self:_get_max_line_end() then
return
end
-- Build fresh sorted ranges so we can search for the first entry at
-- or after from_line, then iterate only the suffix that needs shifting.
self:_rebuild_ranges()
local shifted = false
local msg_ranges = self._message_ranges
local first_msg = first_range_at_or_after(msg_ranges, from_line)
for i = first_msg, #msg_ranges do
local msg_data = self._messages[msg_ranges[i][3]]
if msg_data then
msg_data.line_start = msg_data.line_start + delta
msg_data.line_end = msg_data.line_end + delta
shifted = true
end
end
local part_ranges = self._part_ranges
local first_part = first_range_at_or_after(part_ranges, from_line)
for i = first_part, #part_ranges do
local part_data = self._parts[part_ranges[i][3]]
if part_data then
part_data.line_start = part_data.line_start + delta
part_data.line_end = part_data.line_end + delta
shifted = true
for _, action in ipairs(part_data.actions) do
shift_action(action, delta)
end
end
end
if shifted then
self._ranges_valid = false
if self._max_line_end_valid then
self._max_line_end = self._max_line_end + delta
end
end
end
return RenderState