buildMessageLookups function

MessageLookups buildMessageLookups(List<Message> normalizedMessages, List<Message> messages)

Build pre-computed lookups for O(1) access to message relationships.

Implementation

MessageLookups buildMessageLookups(
  List<Message> normalizedMessages,
  List<Message> messages,
) {
  // Group assistant messages by ID and collect all tool use IDs per message
  final toolUseIDsByMessageID = <String, Set<String>>{};
  final toolUseIDToMessageID = <String, String>{};
  final toolUseByToolUseID = <String, ToolUseBlock>{};

  for (final msg in messages) {
    if (msg is AssistantMessage) {
      final id = msg.messageId;
      toolUseIDsByMessageID.putIfAbsent(id, () => <String>{});
      for (final block in msg.content) {
        if (block is ToolUseBlock) {
          toolUseIDsByMessageID[id]!.add(block.id);
          toolUseIDToMessageID[block.id] = id;
          toolUseByToolUseID[block.id] = block;
        }
      }
    }
  }

  // Build sibling lookup
  final siblingToolUseIDs = <String, Set<String>>{};
  for (final entry in toolUseIDToMessageID.entries) {
    siblingToolUseIDs[entry.key] =
        toolUseIDsByMessageID[entry.value] ?? <String>{};
  }

  // Single pass for progress, hook, and tool result lookups
  final progressMessagesByToolUseID = <String, List<ProgressMessage>>{};
  final inProgressHookCounts = <String, Map<HookEvent, int>>{};
  final resolvedHookNames = <String, Map<HookEvent, Set<String>>>{};
  final toolResultByToolUseID = <String, NormalizedMessage>{};
  final resolvedToolUseIDs = <String>{};
  final erroredToolUseIDs = <String>{};

  for (final msg in normalizedMessages) {
    if (msg is ProgressMessage) {
      progressMessagesByToolUseID
          .putIfAbsent(msg.parentToolUseID, () => [])
          .add(msg);

      if (msg.data['type'] == 'hook_progress') {
        final hookEvent = _parseHookEvent(msg.data['hookEvent']);
        if (hookEvent != null) {
          inProgressHookCounts
              .putIfAbsent(msg.parentToolUseID, () => {})
              .update(hookEvent, (v) => v + 1, ifAbsent: () => 1);
        }
      }
    }

    if (msg is UserMessage) {
      final content = msg.content;
      if (content is List) {
        for (final block in content) {
          if (block is ToolResultBlock) {
            toolResultByToolUseID[block.toolUseId] = NormalizedMessage(msg);
            resolvedToolUseIDs.add(block.toolUseId);
            if (block.isError) erroredToolUseIDs.add(block.toolUseId);
          }
        }
      }
    }
  }

  // Convert resolved hook name sets to counts
  final resolvedHookCounts = <String, Map<HookEvent, int>>{};
  for (final entry in resolvedHookNames.entries) {
    final countMap = <HookEvent, int>{};
    for (final hookEntry in entry.value.entries) {
      countMap[hookEntry.key] = hookEntry.value.length;
    }
    resolvedHookCounts[entry.key] = countMap;
  }

  return MessageLookups(
    siblingToolUseIDs: siblingToolUseIDs,
    progressMessagesByToolUseID: progressMessagesByToolUseID,
    inProgressHookCounts: inProgressHookCounts,
    resolvedHookCounts: resolvedHookCounts,
    toolResultByToolUseID: toolResultByToolUseID,
    toolUseByToolUseID: toolUseByToolUseID,
    normalizedMessageCount: normalizedMessages.length,
    resolvedToolUseIDs: resolvedToolUseIDs,
    erroredToolUseIDs: erroredToolUseIDs,
  );
}