proto.rs

  1#![allow(non_snake_case)]
  2
  3pub mod error;
  4mod macros;
  5mod typed_envelope;
  6
  7pub use error::*;
  8pub use typed_envelope::*;
  9
 10use anyhow::anyhow;
 11use collections::HashMap;
 12use futures::{future::BoxFuture, Future};
 13pub use prost::{DecodeError, Message};
 14use serde::Serialize;
 15use std::{
 16    any::{Any, TypeId},
 17    cmp,
 18    fmt::{self, Debug},
 19    iter, mem,
 20    sync::Arc,
 21    time::{Duration, SystemTime, UNIX_EPOCH},
 22};
 23
 24include!(concat!(env!("OUT_DIR"), "/zed.messages.rs"));
 25
 26pub trait EnvelopedMessage: Clone + Debug + Serialize + Sized + Send + Sync + 'static {
 27    const NAME: &'static str;
 28    const PRIORITY: MessagePriority;
 29    fn into_envelope(
 30        self,
 31        id: u32,
 32        responding_to: Option<u32>,
 33        original_sender_id: Option<PeerId>,
 34    ) -> Envelope;
 35    fn from_envelope(envelope: Envelope) -> Option<Self>;
 36}
 37
 38pub trait EntityMessage: EnvelopedMessage {
 39    type Entity;
 40    fn remote_entity_id(&self) -> u64;
 41}
 42
 43pub trait RequestMessage: EnvelopedMessage {
 44    type Response: EnvelopedMessage;
 45}
 46
 47pub trait AnyTypedEnvelope: 'static + Send + Sync {
 48    fn payload_type_id(&self) -> TypeId;
 49    fn payload_type_name(&self) -> &'static str;
 50    fn as_any(&self) -> &dyn Any;
 51    fn into_any(self: Box<Self>) -> Box<dyn Any + Send + Sync>;
 52    fn is_background(&self) -> bool;
 53    fn original_sender_id(&self) -> Option<PeerId>;
 54    fn sender_id(&self) -> PeerId;
 55    fn message_id(&self) -> u32;
 56}
 57
 58pub enum MessagePriority {
 59    Foreground,
 60    Background,
 61}
 62
 63pub trait ProtoClient: Send + Sync {
 64    fn request(
 65        &self,
 66        envelope: Envelope,
 67        request_type: &'static str,
 68    ) -> BoxFuture<'static, anyhow::Result<Envelope>>;
 69
 70    fn send(&self, envelope: Envelope) -> anyhow::Result<()>;
 71}
 72
 73#[derive(Clone)]
 74pub struct AnyProtoClient(Arc<dyn ProtoClient>);
 75
 76impl<T> From<Arc<T>> for AnyProtoClient
 77where
 78    T: ProtoClient + 'static,
 79{
 80    fn from(client: Arc<T>) -> Self {
 81        Self(client)
 82    }
 83}
 84
 85impl AnyProtoClient {
 86    pub fn new<T: ProtoClient + 'static>(client: Arc<T>) -> Self {
 87        Self(client)
 88    }
 89
 90    pub fn request<T: RequestMessage>(
 91        &self,
 92        request: T,
 93    ) -> impl Future<Output = anyhow::Result<T::Response>> {
 94        let envelope = request.into_envelope(0, None, None);
 95        let response = self.0.request(envelope, T::NAME);
 96        async move {
 97            T::Response::from_envelope(response.await?)
 98                .ok_or_else(|| anyhow!("received response of the wrong type"))
 99        }
100    }
101
102    pub fn send<T: EnvelopedMessage>(&self, request: T) -> anyhow::Result<()> {
103        let envelope = request.into_envelope(0, None, None);
104        self.0.send(envelope)
105    }
106
107    pub fn send_dynamic(&self, message: Envelope) -> anyhow::Result<()> {
108        self.0.send(message)
109    }
110}
111
112impl<T: EnvelopedMessage> AnyTypedEnvelope for TypedEnvelope<T> {
113    fn payload_type_id(&self) -> TypeId {
114        TypeId::of::<T>()
115    }
116
117    fn payload_type_name(&self) -> &'static str {
118        T::NAME
119    }
120
121    fn as_any(&self) -> &dyn Any {
122        self
123    }
124
125    fn into_any(self: Box<Self>) -> Box<dyn Any + Send + Sync> {
126        self
127    }
128
129    fn is_background(&self) -> bool {
130        matches!(T::PRIORITY, MessagePriority::Background)
131    }
132
133    fn original_sender_id(&self) -> Option<PeerId> {
134        self.original_sender_id
135    }
136
137    fn sender_id(&self) -> PeerId {
138        self.sender_id
139    }
140
141    fn message_id(&self) -> u32 {
142        self.message_id
143    }
144}
145
146impl PeerId {
147    pub fn from_u64(peer_id: u64) -> Self {
148        let owner_id = (peer_id >> 32) as u32;
149        let id = peer_id as u32;
150        Self { owner_id, id }
151    }
152
153    pub fn as_u64(self) -> u64 {
154        ((self.owner_id as u64) << 32) | (self.id as u64)
155    }
156}
157
158impl Copy for PeerId {}
159
160impl Eq for PeerId {}
161
162impl Ord for PeerId {
163    fn cmp(&self, other: &Self) -> cmp::Ordering {
164        self.owner_id
165            .cmp(&other.owner_id)
166            .then_with(|| self.id.cmp(&other.id))
167    }
168}
169
170impl PartialOrd for PeerId {
171    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
172        Some(self.cmp(other))
173    }
174}
175
176impl std::hash::Hash for PeerId {
177    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
178        self.owner_id.hash(state);
179        self.id.hash(state);
180    }
181}
182
183impl fmt::Display for PeerId {
184    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
185        write!(f, "{}/{}", self.owner_id, self.id)
186    }
187}
188
189messages!(
190    (Ack, Foreground),
191    (AckBufferOperation, Background),
192    (AckChannelMessage, Background),
193    (AddNotification, Foreground),
194    (AddProjectCollaborator, Foreground),
195    (ApplyCodeAction, Background),
196    (ApplyCodeActionResponse, Background),
197    (ApplyCompletionAdditionalEdits, Background),
198    (ApplyCompletionAdditionalEditsResponse, Background),
199    (BufferReloaded, Foreground),
200    (BufferSaved, Foreground),
201    (Call, Foreground),
202    (CallCanceled, Foreground),
203    (CancelCall, Foreground),
204    (ChannelMessageSent, Foreground),
205    (ChannelMessageUpdate, Foreground),
206    (CompleteWithLanguageModel, Background),
207    (ComputeEmbeddings, Background),
208    (ComputeEmbeddingsResponse, Background),
209    (CopyProjectEntry, Foreground),
210    (CountTokensWithLanguageModel, Background),
211    (CountTokensResponse, Background),
212    (CreateBufferForPeer, Foreground),
213    (CreateChannel, Foreground),
214    (CreateChannelResponse, Foreground),
215    (CreateProjectEntry, Foreground),
216    (CreateRoom, Foreground),
217    (CreateRoomResponse, Foreground),
218    (DeclineCall, Foreground),
219    (DeleteChannel, Foreground),
220    (DeleteNotification, Foreground),
221    (UpdateNotification, Foreground),
222    (DeleteProjectEntry, Foreground),
223    (EndStream, Foreground),
224    (Error, Foreground),
225    (ExpandProjectEntry, Foreground),
226    (ExpandProjectEntryResponse, Foreground),
227    (Follow, Foreground),
228    (FollowResponse, Foreground),
229    (FormatBuffers, Foreground),
230    (FormatBuffersResponse, Foreground),
231    (FuzzySearchUsers, Foreground),
232    (GetCachedEmbeddings, Background),
233    (GetCachedEmbeddingsResponse, Background),
234    (GetChannelMembers, Foreground),
235    (GetChannelMembersResponse, Foreground),
236    (GetChannelMessages, Background),
237    (GetChannelMessagesById, Background),
238    (GetChannelMessagesResponse, Background),
239    (GetCodeActions, Background),
240    (GetCodeActionsResponse, Background),
241    (GetCompletions, Background),
242    (GetCompletionsResponse, Background),
243    (GetDefinition, Background),
244    (GetDefinitionResponse, Background),
245    (GetDocumentHighlights, Background),
246    (GetDocumentHighlightsResponse, Background),
247    (GetHover, Background),
248    (GetHoverResponse, Background),
249    (GetNotifications, Foreground),
250    (GetNotificationsResponse, Foreground),
251    (GetPrivateUserInfo, Foreground),
252    (GetPrivateUserInfoResponse, Foreground),
253    (GetProjectSymbols, Background),
254    (GetProjectSymbolsResponse, Background),
255    (GetReferences, Background),
256    (GetReferencesResponse, Background),
257    (GetSignatureHelp, Background),
258    (GetSignatureHelpResponse, Background),
259    (GetSupermavenApiKey, Background),
260    (GetSupermavenApiKeyResponse, Background),
261    (GetTypeDefinition, Background),
262    (GetTypeDefinitionResponse, Background),
263    (GetImplementation, Background),
264    (GetImplementationResponse, Background),
265    (GetUsers, Foreground),
266    (Hello, Foreground),
267    (IncomingCall, Foreground),
268    (InlayHints, Background),
269    (InlayHintsResponse, Background),
270    (InviteChannelMember, Foreground),
271    (JoinChannel, Foreground),
272    (JoinChannelBuffer, Foreground),
273    (JoinChannelBufferResponse, Foreground),
274    (JoinChannelChat, Foreground),
275    (JoinChannelChatResponse, Foreground),
276    (JoinProject, Foreground),
277    (JoinHostedProject, Foreground),
278    (JoinProjectResponse, Foreground),
279    (JoinRoom, Foreground),
280    (JoinRoomResponse, Foreground),
281    (LanguageModelResponse, Background),
282    (LeaveChannelBuffer, Background),
283    (LeaveChannelChat, Foreground),
284    (LeaveProject, Foreground),
285    (LeaveRoom, Foreground),
286    (MarkNotificationRead, Foreground),
287    (MoveChannel, Foreground),
288    (OnTypeFormatting, Background),
289    (OnTypeFormattingResponse, Background),
290    (OpenBufferById, Background),
291    (OpenBufferByPath, Background),
292    (OpenBufferForSymbol, Background),
293    (OpenBufferForSymbolResponse, Background),
294    (OpenBufferResponse, Background),
295    (PerformRename, Background),
296    (PerformRenameResponse, Background),
297    (Ping, Foreground),
298    (PrepareRename, Background),
299    (PrepareRenameResponse, Background),
300    (ProjectEntryResponse, Foreground),
301    (RefreshInlayHints, Foreground),
302    (RejoinChannelBuffers, Foreground),
303    (RejoinChannelBuffersResponse, Foreground),
304    (RejoinRoom, Foreground),
305    (RejoinRoomResponse, Foreground),
306    (ReloadBuffers, Foreground),
307    (ReloadBuffersResponse, Foreground),
308    (RemoveChannelMember, Foreground),
309    (RemoveChannelMessage, Foreground),
310    (UpdateChannelMessage, Foreground),
311    (RemoveContact, Foreground),
312    (RemoveProjectCollaborator, Foreground),
313    (RenameChannel, Foreground),
314    (RenameChannelResponse, Foreground),
315    (RenameProjectEntry, Foreground),
316    (RequestContact, Foreground),
317    (ResolveCompletionDocumentation, Background),
318    (ResolveCompletionDocumentationResponse, Background),
319    (ResolveInlayHint, Background),
320    (ResolveInlayHintResponse, Background),
321    (RespondToChannelInvite, Foreground),
322    (RespondToContactRequest, Foreground),
323    (RoomUpdated, Foreground),
324    (SaveBuffer, Foreground),
325    (SetChannelMemberRole, Foreground),
326    (SetChannelVisibility, Foreground),
327    (SearchProject, Background),
328    (SearchProjectResponse, Background),
329    (SendChannelMessage, Background),
330    (SendChannelMessageResponse, Background),
331    (ShareProject, Foreground),
332    (ShareProjectResponse, Foreground),
333    (ShowContacts, Foreground),
334    (StartLanguageServer, Foreground),
335    (SubscribeToChannels, Foreground),
336    (SynchronizeBuffers, Foreground),
337    (SynchronizeBuffersResponse, Foreground),
338    (TaskContextForLocation, Background),
339    (TaskContext, Background),
340    (TaskTemplates, Background),
341    (TaskTemplatesResponse, Background),
342    (Test, Foreground),
343    (Unfollow, Foreground),
344    (UnshareProject, Foreground),
345    (UpdateBuffer, Foreground),
346    (UpdateBufferFile, Foreground),
347    (UpdateChannelBuffer, Foreground),
348    (UpdateChannelBufferCollaborators, Foreground),
349    (UpdateChannels, Foreground),
350    (UpdateUserChannels, Foreground),
351    (UpdateContacts, Foreground),
352    (UpdateDiagnosticSummary, Foreground),
353    (UpdateDiffBase, Foreground),
354    (UpdateFollowers, Foreground),
355    (UpdateInviteInfo, Foreground),
356    (UpdateLanguageServer, Foreground),
357    (UpdateParticipantLocation, Foreground),
358    (UpdateProject, Foreground),
359    (UpdateProjectCollaborator, Foreground),
360    (UpdateWorktree, Foreground),
361    (UpdateWorktreeSettings, Foreground),
362    (UsersResponse, Foreground),
363    (LspExtExpandMacro, Background),
364    (LspExtExpandMacroResponse, Background),
365    (SetRoomParticipantRole, Foreground),
366    (BlameBuffer, Foreground),
367    (BlameBufferResponse, Foreground),
368    (CreateDevServerProject, Background),
369    (CreateDevServerProjectResponse, Foreground),
370    (CreateDevServer, Foreground),
371    (CreateDevServerResponse, Foreground),
372    (DevServerInstructions, Foreground),
373    (ShutdownDevServer, Foreground),
374    (ReconnectDevServer, Foreground),
375    (ReconnectDevServerResponse, Foreground),
376    (ShareDevServerProject, Foreground),
377    (JoinDevServerProject, Foreground),
378    (RejoinRemoteProjects, Foreground),
379    (RejoinRemoteProjectsResponse, Foreground),
380    (MultiLspQuery, Background),
381    (MultiLspQueryResponse, Background),
382    (DevServerProjectsUpdate, Foreground),
383    (ValidateDevServerProjectRequest, Background),
384    (ListRemoteDirectory, Background),
385    (ListRemoteDirectoryResponse, Background),
386    (UpdateDevServerProject, Background),
387    (DeleteDevServer, Foreground),
388    (DeleteDevServerProject, Foreground),
389    (RegenerateDevServerToken, Foreground),
390    (RegenerateDevServerTokenResponse, Foreground),
391    (RenameDevServer, Foreground),
392    (OpenNewBuffer, Foreground),
393    (RestartLanguageServers, Foreground),
394    (LinkedEditingRange, Background),
395    (LinkedEditingRangeResponse, Background),
396    (AdvertiseContexts, Foreground),
397    (OpenContext, Foreground),
398    (OpenContextResponse, Foreground),
399    (UpdateContext, Foreground),
400    (SynchronizeContexts, Foreground),
401    (SynchronizeContextsResponse, Foreground),
402    (AddWorktree, Foreground),
403    (AddWorktreeResponse, Foreground),
404);
405
406request_messages!(
407    (ApplyCodeAction, ApplyCodeActionResponse),
408    (
409        ApplyCompletionAdditionalEdits,
410        ApplyCompletionAdditionalEditsResponse
411    ),
412    (Call, Ack),
413    (CancelCall, Ack),
414    (CopyProjectEntry, ProjectEntryResponse),
415    (CompleteWithLanguageModel, LanguageModelResponse),
416    (ComputeEmbeddings, ComputeEmbeddingsResponse),
417    (CountTokensWithLanguageModel, CountTokensResponse),
418    (CreateChannel, CreateChannelResponse),
419    (CreateProjectEntry, ProjectEntryResponse),
420    (CreateRoom, CreateRoomResponse),
421    (DeclineCall, Ack),
422    (DeleteChannel, Ack),
423    (DeleteProjectEntry, ProjectEntryResponse),
424    (ExpandProjectEntry, ExpandProjectEntryResponse),
425    (Follow, FollowResponse),
426    (FormatBuffers, FormatBuffersResponse),
427    (FuzzySearchUsers, UsersResponse),
428    (GetCachedEmbeddings, GetCachedEmbeddingsResponse),
429    (GetChannelMembers, GetChannelMembersResponse),
430    (GetChannelMessages, GetChannelMessagesResponse),
431    (GetChannelMessagesById, GetChannelMessagesResponse),
432    (GetCodeActions, GetCodeActionsResponse),
433    (GetCompletions, GetCompletionsResponse),
434    (GetDefinition, GetDefinitionResponse),
435    (GetImplementation, GetImplementationResponse),
436    (GetDocumentHighlights, GetDocumentHighlightsResponse),
437    (GetHover, GetHoverResponse),
438    (GetNotifications, GetNotificationsResponse),
439    (GetPrivateUserInfo, GetPrivateUserInfoResponse),
440    (GetProjectSymbols, GetProjectSymbolsResponse),
441    (GetReferences, GetReferencesResponse),
442    (GetSignatureHelp, GetSignatureHelpResponse),
443    (GetSupermavenApiKey, GetSupermavenApiKeyResponse),
444    (GetTypeDefinition, GetTypeDefinitionResponse),
445    (LinkedEditingRange, LinkedEditingRangeResponse),
446    (ListRemoteDirectory, ListRemoteDirectoryResponse),
447    (UpdateDevServerProject, Ack),
448    (GetUsers, UsersResponse),
449    (IncomingCall, Ack),
450    (InlayHints, InlayHintsResponse),
451    (InviteChannelMember, Ack),
452    (JoinChannel, JoinRoomResponse),
453    (JoinChannelBuffer, JoinChannelBufferResponse),
454    (JoinChannelChat, JoinChannelChatResponse),
455    (JoinHostedProject, JoinProjectResponse),
456    (JoinProject, JoinProjectResponse),
457    (JoinRoom, JoinRoomResponse),
458    (LeaveChannelBuffer, Ack),
459    (LeaveRoom, Ack),
460    (MarkNotificationRead, Ack),
461    (MoveChannel, Ack),
462    (OnTypeFormatting, OnTypeFormattingResponse),
463    (OpenBufferById, OpenBufferResponse),
464    (OpenBufferByPath, OpenBufferResponse),
465    (OpenBufferForSymbol, OpenBufferForSymbolResponse),
466    (OpenNewBuffer, OpenBufferResponse),
467    (PerformRename, PerformRenameResponse),
468    (Ping, Ack),
469    (PrepareRename, PrepareRenameResponse),
470    (RefreshInlayHints, Ack),
471    (RejoinChannelBuffers, RejoinChannelBuffersResponse),
472    (RejoinRoom, RejoinRoomResponse),
473    (ReloadBuffers, ReloadBuffersResponse),
474    (RemoveChannelMember, Ack),
475    (RemoveChannelMessage, Ack),
476    (UpdateChannelMessage, Ack),
477    (RemoveContact, Ack),
478    (RenameChannel, RenameChannelResponse),
479    (RenameProjectEntry, ProjectEntryResponse),
480    (RequestContact, Ack),
481    (
482        ResolveCompletionDocumentation,
483        ResolveCompletionDocumentationResponse
484    ),
485    (ResolveInlayHint, ResolveInlayHintResponse),
486    (RespondToChannelInvite, Ack),
487    (RespondToContactRequest, Ack),
488    (SaveBuffer, BufferSaved),
489    (SearchProject, SearchProjectResponse),
490    (SendChannelMessage, SendChannelMessageResponse),
491    (SetChannelMemberRole, Ack),
492    (SetChannelVisibility, Ack),
493    (ShareProject, ShareProjectResponse),
494    (SynchronizeBuffers, SynchronizeBuffersResponse),
495    (TaskContextForLocation, TaskContext),
496    (TaskTemplates, TaskTemplatesResponse),
497    (Test, Test),
498    (UpdateBuffer, Ack),
499    (UpdateParticipantLocation, Ack),
500    (UpdateProject, Ack),
501    (UpdateWorktree, Ack),
502    (LspExtExpandMacro, LspExtExpandMacroResponse),
503    (SetRoomParticipantRole, Ack),
504    (BlameBuffer, BlameBufferResponse),
505    (CreateDevServerProject, CreateDevServerProjectResponse),
506    (CreateDevServer, CreateDevServerResponse),
507    (ShutdownDevServer, Ack),
508    (ShareDevServerProject, ShareProjectResponse),
509    (JoinDevServerProject, JoinProjectResponse),
510    (RejoinRemoteProjects, RejoinRemoteProjectsResponse),
511    (ReconnectDevServer, ReconnectDevServerResponse),
512    (ValidateDevServerProjectRequest, Ack),
513    (MultiLspQuery, MultiLspQueryResponse),
514    (DeleteDevServer, Ack),
515    (DeleteDevServerProject, Ack),
516    (RegenerateDevServerToken, RegenerateDevServerTokenResponse),
517    (RenameDevServer, Ack),
518    (RestartLanguageServers, Ack),
519    (OpenContext, OpenContextResponse),
520    (SynchronizeContexts, SynchronizeContextsResponse),
521    (AddWorktree, AddWorktreeResponse),
522);
523
524entity_messages!(
525    {project_id, ShareProject},
526    AddProjectCollaborator,
527    ApplyCodeAction,
528    ApplyCompletionAdditionalEdits,
529    BlameBuffer,
530    BufferReloaded,
531    BufferSaved,
532    CopyProjectEntry,
533    CreateBufferForPeer,
534    CreateProjectEntry,
535    DeleteProjectEntry,
536    ExpandProjectEntry,
537    FormatBuffers,
538    GetCodeActions,
539    GetCompletions,
540    GetDefinition,
541    GetImplementation,
542    GetDocumentHighlights,
543    GetHover,
544    GetProjectSymbols,
545    GetReferences,
546    GetSignatureHelp,
547    GetTypeDefinition,
548    InlayHints,
549    JoinProject,
550    LeaveProject,
551    LinkedEditingRange,
552    MultiLspQuery,
553    RestartLanguageServers,
554    OnTypeFormatting,
555    OpenNewBuffer,
556    OpenBufferById,
557    OpenBufferByPath,
558    OpenBufferForSymbol,
559    PerformRename,
560    PrepareRename,
561    RefreshInlayHints,
562    ReloadBuffers,
563    RemoveProjectCollaborator,
564    RenameProjectEntry,
565    ResolveCompletionDocumentation,
566    ResolveInlayHint,
567    SaveBuffer,
568    SearchProject,
569    StartLanguageServer,
570    SynchronizeBuffers,
571    TaskContextForLocation,
572    TaskTemplates,
573    UnshareProject,
574    UpdateBuffer,
575    UpdateBufferFile,
576    UpdateDiagnosticSummary,
577    UpdateDiffBase,
578    UpdateLanguageServer,
579    UpdateProject,
580    UpdateProjectCollaborator,
581    UpdateWorktree,
582    UpdateWorktreeSettings,
583    LspExtExpandMacro,
584    AdvertiseContexts,
585    OpenContext,
586    UpdateContext,
587    SynchronizeContexts,
588);
589
590entity_messages!(
591    {channel_id, Channel},
592    ChannelMessageSent,
593    ChannelMessageUpdate,
594    RemoveChannelMessage,
595    UpdateChannelMessage,
596    UpdateChannelBuffer,
597    UpdateChannelBufferCollaborators,
598);
599
600impl From<Timestamp> for SystemTime {
601    fn from(val: Timestamp) -> Self {
602        UNIX_EPOCH
603            .checked_add(Duration::new(val.seconds, val.nanos))
604            .unwrap()
605    }
606}
607
608impl From<SystemTime> for Timestamp {
609    fn from(time: SystemTime) -> Self {
610        let duration = time.duration_since(UNIX_EPOCH).unwrap();
611        Self {
612            seconds: duration.as_secs(),
613            nanos: duration.subsec_nanos(),
614        }
615    }
616}
617
618impl From<u128> for Nonce {
619    fn from(nonce: u128) -> Self {
620        let upper_half = (nonce >> 64) as u64;
621        let lower_half = nonce as u64;
622        Self {
623            upper_half,
624            lower_half,
625        }
626    }
627}
628
629impl From<Nonce> for u128 {
630    fn from(nonce: Nonce) -> Self {
631        let upper_half = (nonce.upper_half as u128) << 64;
632        let lower_half = nonce.lower_half as u128;
633        upper_half | lower_half
634    }
635}
636
637pub fn split_worktree_update(
638    mut message: UpdateWorktree,
639    max_chunk_size: usize,
640) -> impl Iterator<Item = UpdateWorktree> {
641    let mut done_files = false;
642
643    let mut repository_map = message
644        .updated_repositories
645        .into_iter()
646        .map(|repo| (repo.work_directory_id, repo))
647        .collect::<HashMap<_, _>>();
648
649    iter::from_fn(move || {
650        if done_files {
651            return None;
652        }
653
654        let updated_entries_chunk_size = cmp::min(message.updated_entries.len(), max_chunk_size);
655        let updated_entries: Vec<_> = message
656            .updated_entries
657            .drain(..updated_entries_chunk_size)
658            .collect();
659
660        let removed_entries_chunk_size = cmp::min(message.removed_entries.len(), max_chunk_size);
661        let removed_entries = message
662            .removed_entries
663            .drain(..removed_entries_chunk_size)
664            .collect();
665
666        done_files = message.updated_entries.is_empty() && message.removed_entries.is_empty();
667
668        let mut updated_repositories = Vec::new();
669
670        if !repository_map.is_empty() {
671            for entry in &updated_entries {
672                if let Some(repo) = repository_map.remove(&entry.id) {
673                    updated_repositories.push(repo)
674                }
675            }
676        }
677
678        let removed_repositories = if done_files {
679            mem::take(&mut message.removed_repositories)
680        } else {
681            Default::default()
682        };
683
684        if done_files {
685            updated_repositories.extend(mem::take(&mut repository_map).into_values());
686        }
687
688        Some(UpdateWorktree {
689            project_id: message.project_id,
690            worktree_id: message.worktree_id,
691            root_name: message.root_name.clone(),
692            abs_path: message.abs_path.clone(),
693            updated_entries,
694            removed_entries,
695            scan_id: message.scan_id,
696            is_last_update: done_files && message.is_last_update,
697            updated_repositories,
698            removed_repositories,
699        })
700    })
701}
702
703#[cfg(test)]
704mod tests {
705    use super::*;
706
707    #[test]
708    fn test_converting_peer_id_from_and_to_u64() {
709        let peer_id = PeerId {
710            owner_id: 10,
711            id: 3,
712        };
713        assert_eq!(PeerId::from_u64(peer_id.as_u64()), peer_id);
714        let peer_id = PeerId {
715            owner_id: u32::MAX,
716            id: 3,
717        };
718        assert_eq!(PeerId::from_u64(peer_id.as_u64()), peer_id);
719        let peer_id = PeerId {
720            owner_id: 10,
721            id: u32::MAX,
722        };
723        assert_eq!(PeerId::from_u64(peer_id.as_u64()), peer_id);
724        let peer_id = PeerId {
725            owner_id: u32::MAX,
726            id: u32::MAX,
727        };
728        assert_eq!(PeerId::from_u64(peer_id.as_u64()), peer_id);
729    }
730}