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