[Media Router] Assign each route a current cast mode if possible

chrome/browser/ui/webui/media_router/media_router_ui.cc

Index: chrome/browser/ui/webui/media_router/media_router_ui.cc
diff --git a/chrome/browser/ui/webui/media_router/media_router_ui.cc b/chrome/browser/ui/webui/media_router/media_router_ui.cc
index bb045896b2dbbab917e2daab026bbe6732e93690..7bf48da12bc8bf37cdacbe838514c0cf031dd328 100644
--- a/chrome/browser/ui/webui/media_router/media_router_ui.cc
+++ b/chrome/browser/ui/webui/media_router/media_router_ui.cc
@@ -123,9 +123,13 @@ class MediaRouterUI::UIIssuesObserver : public IssuesObserver {
 };
 
 MediaRouterUI::UIMediaRoutesObserver::UIMediaRoutesObserver(
-    MediaRouter* router, const MediaSource::Id& source_id,
+    MediaRouter* router,
+    const MediaSource::Id& source_id,
+    const SourceCastModeMapCallback& source_cast_mode_callback,
     const RoutesUpdatedCallback& callback)
-    : MediaRoutesObserver(router, source_id), callback_(callback) {
+    : MediaRoutesObserver(router, source_id),
+      source_cast_mode_callback_(source_cast_mode_callback),
+      callback_(callback) {
   DCHECK(!callback_.is_null());
 }
 
@@ -156,7 +160,18 @@ void MediaRouterUI::UIMediaRoutesObserver::OnRoutesUpdated(
     }
   }
 
-  callback_.Run(routes_for_display, joinable_route_ids_for_display);
+  auto available_source_map(source_cast_mode_callback_.Run());
+  std::unordered_map<MediaRoute::Id, MediaCastMode> current_cast_modes;
+  for (const auto& route : routes_for_display) {
+    auto source_entry = available_source_map.find(route.media_source().id());
+    if (source_entry != end(available_source_map)) {
+      current_cast_modes.insert(
+          std::make_pair(route.media_route_id(), source_entry->second));
+    }
+  }
+
+  callback_.Run(routes_for_display, joinable_route_ids_for_display,
+                current_cast_modes);
 }
 
 MediaRouterUI::MediaRouterUI(content::WebUI* web_ui)
@@ -224,8 +239,11 @@ void MediaRouterUI::InitWithDefaultMediaSource(
         presentation_service_delegate_->GetDefaultPresentationRequest());
   } else {
     // Register for MediaRoute updates without a media source.
-    routes_observer_.reset(new UIMediaRoutesObserver(router_, MediaSource::Id(),
-          base::Bind(&MediaRouterUI::OnRoutesUpdated, base::Unretained(this))));
+    routes_observer_.reset(new UIMediaRoutesObserver(
+        router_, MediaSource::Id(),
+        base::Bind(&MediaRouterUI::GetAvailableSourceMap,
+                   base::Unretained(this)),
+        base::Bind(&MediaRouterUI::OnRoutesUpdated, base::Unretained(this))));
   }
 }
 
@@ -302,6 +320,7 @@ void MediaRouterUI::OnDefaultPresentationChanged(
   // Register for MediaRoute updates.
   routes_observer_.reset(new UIMediaRoutesObserver(
       router_, source.id(),
+      base::Bind(&MediaRouterUI::GetAvailableSourceMap, base::Unretained(this)),
       base::Bind(&MediaRouterUI::OnRoutesUpdated, base::Unretained(this))));
 
   UpdateCastModes();
@@ -312,8 +331,9 @@ void MediaRouterUI::OnDefaultPresentationRemoved() {
   query_result_manager_->StopSinksQuery(MediaCastMode::DEFAULT);
   // Register for MediaRoute updates without a media source.
   routes_observer_.reset(new UIMediaRoutesObserver(
-        router_, MediaSource::Id(),
-        base::Bind(&MediaRouterUI::OnRoutesUpdated, base::Unretained(this))));
+      router_, MediaSource::Id(),
+      base::Bind(&MediaRouterUI::GetAvailableSourceMap, base::Unretained(this)),
+      base::Bind(&MediaRouterUI::OnRoutesUpdated, base::Unretained(this))));
   UpdateCastModes();
 }
 
@@ -512,12 +532,28 @@ void MediaRouterUI::SetIssue(const Issue* issue) {
   if (ui_initialized_) handler_->UpdateIssue(issue);
 }
 
+std::unordered_map<MediaSource::Id, MediaCastMode>
+MediaRouterUI::GetAvailableSourceMap() {
+  std::unordered_map<MediaSource::Id, MediaCastMode> available_source_map;
+  for (const auto& cast_mode : cast_modes_) {
+    available_source_map.insert(std::make_pair(
+        query_result_manager_->GetSourceForCastMode(cast_mode).id(),
+        cast_mode));
+  }
+  return available_source_map;
+}
+
 void MediaRouterUI::OnRoutesUpdated(
     const std::vector<MediaRoute>& routes,
-    const std::vector<MediaRoute::Id>& joinable_route_ids) {
+    const std::vector<MediaRoute::Id>& joinable_route_ids,
+    const std::unordered_map<MediaRoute::Id, MediaCastMode>&
+        current_cast_modes) {
   routes_ = routes;
   joinable_route_ids_ = joinable_route_ids;
-  if (ui_initialized_) handler_->UpdateRoutes(routes_, joinable_route_ids_);
+  current_cast_modes_ = current_cast_modes;
+
+  if (ui_initialized_)
+    handler_->UpdateRoutes(routes_, joinable_route_ids_, current_cast_modes_);
 }
 
 void MediaRouterUI::OnRouteResponseReceived(