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    // Remote development
399    (AddWorktree, Foreground),
400    (AddWorktreeResponse, Foreground),
401);
402
403request_messages!(
404    (ApplyCodeAction, ApplyCodeActionResponse),
405    (
406        ApplyCompletionAdditionalEdits,
407        ApplyCompletionAdditionalEditsResponse
408    ),
409    (Call, Ack),
410    (CancelCall, Ack),
411    (CopyProjectEntry, ProjectEntryResponse),
412    (CompleteWithLanguageModel, LanguageModelResponse),
413    (ComputeEmbeddings, ComputeEmbeddingsResponse),
414    (CountTokensWithLanguageModel, CountTokensResponse),
415    (CreateChannel, CreateChannelResponse),
416    (CreateProjectEntry, ProjectEntryResponse),
417    (CreateRoom, CreateRoomResponse),
418    (DeclineCall, Ack),
419    (DeleteChannel, Ack),
420    (DeleteProjectEntry, ProjectEntryResponse),
421    (ExpandProjectEntry, ExpandProjectEntryResponse),
422    (Follow, FollowResponse),
423    (FormatBuffers, FormatBuffersResponse),
424    (FuzzySearchUsers, UsersResponse),
425    (GetCachedEmbeddings, GetCachedEmbeddingsResponse),
426    (GetChannelMembers, GetChannelMembersResponse),
427    (GetChannelMessages, GetChannelMessagesResponse),
428    (GetChannelMessagesById, GetChannelMessagesResponse),
429    (GetCodeActions, GetCodeActionsResponse),
430    (GetCompletions, GetCompletionsResponse),
431    (GetDefinition, GetDefinitionResponse),
432    (GetImplementation, GetImplementationResponse),
433    (GetDocumentHighlights, GetDocumentHighlightsResponse),
434    (GetHover, GetHoverResponse),
435    (GetNotifications, GetNotificationsResponse),
436    (GetPrivateUserInfo, GetPrivateUserInfoResponse),
437    (GetProjectSymbols, GetProjectSymbolsResponse),
438    (GetReferences, GetReferencesResponse),
439    (GetSignatureHelp, GetSignatureHelpResponse),
440    (GetSupermavenApiKey, GetSupermavenApiKeyResponse),
441    (GetTypeDefinition, GetTypeDefinitionResponse),
442    (LinkedEditingRange, LinkedEditingRangeResponse),
443    (ListRemoteDirectory, ListRemoteDirectoryResponse),
444    (UpdateDevServerProject, Ack),
445    (GetUsers, UsersResponse),
446    (IncomingCall, Ack),
447    (InlayHints, InlayHintsResponse),
448    (InviteChannelMember, Ack),
449    (JoinChannel, JoinRoomResponse),
450    (JoinChannelBuffer, JoinChannelBufferResponse),
451    (JoinChannelChat, JoinChannelChatResponse),
452    (JoinHostedProject, JoinProjectResponse),
453    (JoinProject, JoinProjectResponse),
454    (JoinRoom, JoinRoomResponse),
455    (LeaveChannelBuffer, Ack),
456    (LeaveRoom, Ack),
457    (MarkNotificationRead, Ack),
458    (MoveChannel, Ack),
459    (OnTypeFormatting, OnTypeFormattingResponse),
460    (OpenBufferById, OpenBufferResponse),
461    (OpenBufferByPath, OpenBufferResponse),
462    (OpenBufferForSymbol, OpenBufferForSymbolResponse),
463    (OpenNewBuffer, OpenBufferResponse),
464    (PerformRename, PerformRenameResponse),
465    (Ping, Ack),
466    (PrepareRename, PrepareRenameResponse),
467    (RefreshInlayHints, Ack),
468    (RejoinChannelBuffers, RejoinChannelBuffersResponse),
469    (RejoinRoom, RejoinRoomResponse),
470    (ReloadBuffers, ReloadBuffersResponse),
471    (RemoveChannelMember, Ack),
472    (RemoveChannelMessage, Ack),
473    (UpdateChannelMessage, Ack),
474    (RemoveContact, Ack),
475    (RenameChannel, RenameChannelResponse),
476    (RenameProjectEntry, ProjectEntryResponse),
477    (RequestContact, Ack),
478    (
479        ResolveCompletionDocumentation,
480        ResolveCompletionDocumentationResponse
481    ),
482    (ResolveInlayHint, ResolveInlayHintResponse),
483    (RespondToChannelInvite, Ack),
484    (RespondToContactRequest, Ack),
485    (SaveBuffer, BufferSaved),
486    (SearchProject, SearchProjectResponse),
487    (SendChannelMessage, SendChannelMessageResponse),
488    (SetChannelMemberRole, Ack),
489    (SetChannelVisibility, Ack),
490    (ShareProject, ShareProjectResponse),
491    (SynchronizeBuffers, SynchronizeBuffersResponse),
492    (TaskContextForLocation, TaskContext),
493    (TaskTemplates, TaskTemplatesResponse),
494    (Test, Test),
495    (UpdateBuffer, Ack),
496    (UpdateParticipantLocation, Ack),
497    (UpdateProject, Ack),
498    (UpdateWorktree, Ack),
499    (LspExtExpandMacro, LspExtExpandMacroResponse),
500    (SetRoomParticipantRole, Ack),
501    (BlameBuffer, BlameBufferResponse),
502    (CreateDevServerProject, CreateDevServerProjectResponse),
503    (CreateDevServer, CreateDevServerResponse),
504    (ShutdownDevServer, Ack),
505    (ShareDevServerProject, ShareProjectResponse),
506    (JoinDevServerProject, JoinProjectResponse),
507    (RejoinRemoteProjects, RejoinRemoteProjectsResponse),
508    (ReconnectDevServer, ReconnectDevServerResponse),
509    (ValidateDevServerProjectRequest, Ack),
510    (MultiLspQuery, MultiLspQueryResponse),
511    (DeleteDevServer, Ack),
512    (DeleteDevServerProject, Ack),
513    (RegenerateDevServerToken, RegenerateDevServerTokenResponse),
514    (RenameDevServer, Ack),
515    (RestartLanguageServers, Ack),
516    (OpenContext, OpenContextResponse),
517    (SynchronizeContexts, SynchronizeContextsResponse),
518    // Remote development
519    (AddWorktree, AddWorktreeResponse),
520);
521
522entity_messages!(
523    {project_id, ShareProject},
524    AddProjectCollaborator,
525    ApplyCodeAction,
526    ApplyCompletionAdditionalEdits,
527    BlameBuffer,
528    BufferReloaded,
529    BufferSaved,
530    CopyProjectEntry,
531    CreateBufferForPeer,
532    CreateProjectEntry,
533    DeleteProjectEntry,
534    ExpandProjectEntry,
535    FormatBuffers,
536    GetCodeActions,
537    GetCompletions,
538    GetDefinition,
539    GetImplementation,
540    GetDocumentHighlights,
541    GetHover,
542    GetProjectSymbols,
543    GetReferences,
544    GetSignatureHelp,
545    GetTypeDefinition,
546    InlayHints,
547    JoinProject,
548    LeaveProject,
549    LinkedEditingRange,
550    MultiLspQuery,
551    RestartLanguageServers,
552    OnTypeFormatting,
553    OpenNewBuffer,
554    OpenBufferById,
555    OpenBufferByPath,
556    OpenBufferForSymbol,
557    PerformRename,
558    PrepareRename,
559    RefreshInlayHints,
560    ReloadBuffers,
561    RemoveProjectCollaborator,
562    RenameProjectEntry,
563    ResolveCompletionDocumentation,
564    ResolveInlayHint,
565    SaveBuffer,
566    SearchProject,
567    StartLanguageServer,
568    SynchronizeBuffers,
569    TaskContextForLocation,
570    TaskTemplates,
571    UnshareProject,
572    UpdateBuffer,
573    UpdateBufferFile,
574    UpdateDiagnosticSummary,
575    UpdateDiffBase,
576    UpdateLanguageServer,
577    UpdateProject,
578    UpdateProjectCollaborator,
579    UpdateWorktree,
580    UpdateWorktreeSettings,
581    LspExtExpandMacro,
582    AdvertiseContexts,
583    OpenContext,
584    UpdateContext,
585    SynchronizeContexts,
586);
587
588entity_messages!(
589    {channel_id, Channel},
590    ChannelMessageSent,
591    ChannelMessageUpdate,
592    RemoveChannelMessage,
593    UpdateChannelMessage,
594    UpdateChannelBuffer,
595    UpdateChannelBufferCollaborators,
596);
597
598impl From<Timestamp> for SystemTime {
599    fn from(val: Timestamp) -> Self {
600        UNIX_EPOCH
601            .checked_add(Duration::new(val.seconds, val.nanos))
602            .unwrap()
603    }
604}
605
606impl From<SystemTime> for Timestamp {
607    fn from(time: SystemTime) -> Self {
608        let duration = time.duration_since(UNIX_EPOCH).unwrap();
609        Self {
610            seconds: duration.as_secs(),
611            nanos: duration.subsec_nanos(),
612        }
613    }
614}
615
616impl From<u128> for Nonce {
617    fn from(nonce: u128) -> Self {
618        let upper_half = (nonce >> 64) as u64;
619        let lower_half = nonce as u64;
620        Self {
621            upper_half,
622            lower_half,
623        }
624    }
625}
626
627impl From<Nonce> for u128 {
628    fn from(nonce: Nonce) -> Self {
629        let upper_half = (nonce.upper_half as u128) << 64;
630        let lower_half = nonce.lower_half as u128;
631        upper_half | lower_half
632    }
633}
634
635pub fn split_worktree_update(
636    mut message: UpdateWorktree,
637    max_chunk_size: usize,
638) -> impl Iterator<Item = UpdateWorktree> {
639    let mut done_files = false;
640
641    let mut repository_map = message
642        .updated_repositories
643        .into_iter()
644        .map(|repo| (repo.work_directory_id, repo))
645        .collect::<HashMap<_, _>>();
646
647    iter::from_fn(move || {
648        if done_files {
649            return None;
650        }
651
652        let updated_entries_chunk_size = cmp::min(message.updated_entries.len(), max_chunk_size);
653        let updated_entries: Vec<_> = message
654            .updated_entries
655            .drain(..updated_entries_chunk_size)
656            .collect();
657
658        let removed_entries_chunk_size = cmp::min(message.removed_entries.len(), max_chunk_size);
659        let removed_entries = message
660            .removed_entries
661            .drain(..removed_entries_chunk_size)
662            .collect();
663
664        done_files = message.updated_entries.is_empty() && message.removed_entries.is_empty();
665
666        let mut updated_repositories = Vec::new();
667
668        if !repository_map.is_empty() {
669            for entry in &updated_entries {
670                if let Some(repo) = repository_map.remove(&entry.id) {
671                    updated_repositories.push(repo)
672                }
673            }
674        }
675
676        let removed_repositories = if done_files {
677            mem::take(&mut message.removed_repositories)
678        } else {
679            Default::default()
680        };
681
682        if done_files {
683            updated_repositories.extend(mem::take(&mut repository_map).into_values());
684        }
685
686        Some(UpdateWorktree {
687            project_id: message.project_id,
688            worktree_id: message.worktree_id,
689            root_name: message.root_name.clone(),
690            abs_path: message.abs_path.clone(),
691            updated_entries,
692            removed_entries,
693            scan_id: message.scan_id,
694            is_last_update: done_files && message.is_last_update,
695            updated_repositories,
696            removed_repositories,
697        })
698    })
699}
700
701#[cfg(test)]
702mod tests {
703    use super::*;
704
705    #[test]
706    fn test_converting_peer_id_from_and_to_u64() {
707        let peer_id = PeerId {
708            owner_id: 10,
709            id: 3,
710        };
711        assert_eq!(PeerId::from_u64(peer_id.as_u64()), peer_id);
712        let peer_id = PeerId {
713            owner_id: u32::MAX,
714            id: 3,
715        };
716        assert_eq!(PeerId::from_u64(peer_id.as_u64()), peer_id);
717        let peer_id = PeerId {
718            owner_id: 10,
719            id: u32::MAX,
720        };
721        assert_eq!(PeerId::from_u64(peer_id.as_u64()), peer_id);
722        let peer_id = PeerId {
723            owner_id: u32::MAX,
724            id: u32::MAX,
725        };
726        assert_eq!(PeerId::from_u64(peer_id.as_u64()), peer_id);
727    }
728}