Implement RCTFabricSurface.synchronouslyWaitForStage

Summary:
Changelog: [Internal]

Exposes `synchronouslyWaitForStage` to `RCTFabricSurface`.

This is a first step towards having screenshot tests rendered with Fabric.

Reviewed By: shergin

Differential Revision: D19603837

fbshipit-source-id: 26c14cf3bbd67fea96319ff08d3321557ddcdd9c

React/Fabric/RCTScheduler.h

@@ -62,6 +62,8 @@ NS_ASSUME_NONNULL_BEGIN
 - (facebook::react::ComponentDescriptor const *)findComponentDescriptorByHandle_DO_NOT_USE_THIS_IS_BROKEN:
     (facebook::react::ComponentHandle)handle;
 
+- (facebook::react::MountingCoordinator::Shared)mountingCoordinatorWithSurfaceId:(facebook::react::SurfaceId)surfaceId;
+
 @end
 
 NS_ASSUME_NONNULL_END

React/Fabric/RCTScheduler.mm

@@ -122,4 +122,9 @@ - (ComponentDescriptor const *)findComponentDescriptorByHandle_DO_NOT_USE_THIS_I
   return _scheduler->findComponentDescriptorByHandle_DO_NOT_USE_THIS_IS_BROKEN(handle);
 }
 
+- (MountingCoordinator::Shared)mountingCoordinatorWithSurfaceId:(SurfaceId)surfaceId
+{
+  return _scheduler->findMountingCoordinator(surfaceId);
+}
+
 @end

React/Fabric/RCTSurfacePresenter.h

@@ -9,7 +9,7 @@
 
 #import <React/RCTPrimitives.h>
 #import <React/RCTSurfacePresenterStub.h>
-
+#import <React/RCTSurfaceStage.h>
 #import <react/utils/ContextContainer.h>
 #import <react/utils/RuntimeExecutor.h>
 
@@ -69,6 +69,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (BOOL)synchronouslyUpdateViewOnUIThread:(NSNumber *)reactTag props:(NSDictionary *)props;
 
+- (BOOL)synchronouslyWaitSurface:(RCTFabricSurface *)surface timeout:(NSTimeInterval)timeout;
+
 - (void)addObserver:(id<RCTSurfacePresenterObserver>)observer;
 
 - (void)removeObserver:(id<RCTSurfacePresenterObserver>)observer;

React/Fabric/RCTSurfacePresenter.mm

@@ -174,6 +174,19 @@ - (BOOL)synchronouslyUpdateViewOnUIThread:(NSNumber *)reactTag props:(NSDictiona
   return YES;
 }
 
+- (BOOL)synchronouslyWaitSurface:(RCTFabricSurface *)surface timeout:(NSTimeInterval)timeout
+{
+  auto mountingCoordinator = [_scheduler mountingCoordinatorWithSurfaceId:surface.rootTag];
+
+  if (!mountingCoordinator->waitForTransaction(std::chrono::duration<NSTimeInterval>(timeout))) {
+    return NO;
+  }
+
+  [_mountingManager scheduleTransaction:mountingCoordinator];
+
+  return YES;
+}
+
 - (BOOL)suspend
 {
   std::unique_lock<better::shared_mutex> lock(_schedulerMutex);

React/Fabric/Surface/RCTFabricSurface.h

@@ -7,6 +7,7 @@
 
 #import <React/RCTPrimitives.h>
 #import <React/RCTSurfaceStage.h>
+#import <react/mounting/MountingCoordinator.h>
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -121,10 +122,9 @@ NS_ASSUME_NONNULL_BEGIN
 
 /**
  * Synchronously blocks the current thread up to given `timeout` until
- * the Surface reaches `stage`.
- * NOT SUPPORTED IN FABRIC YET.
+ * the Surface is rendered.
  */
-- (BOOL)synchronouslyWaitForStage:(RCTSurfaceStage)stage timeout:(NSTimeInterval)timeout;
+- (BOOL)synchronouslyWaitFor:(NSTimeInterval)timeout;
 
 @end
 

React/Fabric/Surface/RCTFabricSurface.mm

@@ -20,6 +20,8 @@
 
 #import "RCTSurfacePresenter.h"
 
+using namespace facebook::react;
+
 @implementation RCTFabricSurface {
   // Immutable
   RCTSurfacePresenter *_surfacePresenter;
@@ -253,10 +255,9 @@ - (CGSize)intrinsicSize
 
 #pragma mark - Synchronous Waiting
 
-- (BOOL)synchronouslyWaitForStage:(RCTSurfaceStage)stage timeout:(NSTimeInterval)timeout
+- (BOOL)synchronouslyWaitFor:(NSTimeInterval)timeout
 {
-  // TODO: Not supported yet.
-  return NO;
+  return [_surfacePresenter synchronouslyWaitSurface:self timeout:timeout];
 }
 
 #pragma mark - Deprecated

ReactCommon/fabric/mounting/ShadowTree.cpp

@@ -124,6 +124,10 @@ Tag ShadowTree::getSurfaceId() const {
   return surfaceId_;
 }
 
+MountingCoordinator::Shared ShadowTree::getMountingCoordinator() const {
+  return mountingCoordinator_;
+}
+
 void ShadowTree::commit(ShadowTreeCommitTransaction transaction) const {
   SystraceSection s("ShadowTree::commit");
 

ReactCommon/fabric/mounting/ShadowTree.h

@@ -7,7 +7,6 @@
 
 #pragma once
 
-
 #include <better/mutex.h>
 #include <memory>
 
@@ -66,6 +65,8 @@ class ShadowTree final {
    */
   void commitEmptyTree() const;
 
+  MountingCoordinator::Shared getMountingCoordinator() const;
+
  private:
   RootShadowNode::Unshared cloneRootShadowNode(
       RootShadowNode::Shared const &oldRootShadowNode,

ReactCommon/fabric/uimanager/Scheduler.cpp

@@ -242,6 +242,16 @@ Size Scheduler::measureSurface(
   return size;
 }
 
+MountingCoordinator::Shared Scheduler::findMountingCoordinator(
+    SurfaceId surfaceId) const {
+  MountingCoordinator::Shared mountingCoordinator = nullptr;
+  uiManager_->getShadowTreeRegistry().visit(
+      surfaceId, [&](const ShadowTree &shadowTree) {
+        mountingCoordinator = shadowTree.getMountingCoordinator();
+      });
+  return mountingCoordinator;
+}
+
 void Scheduler::constraintSurfaceLayout(
     SurfaceId surfaceId,
     const LayoutConstraints &layoutConstraints,

ReactCommon/fabric/uimanager/Scheduler.h

@@ -75,6 +75,9 @@ class Scheduler final : public UIManagerDelegate {
   findComponentDescriptorByHandle_DO_NOT_USE_THIS_IS_BROKEN(
       ComponentHandle handle) const;
 
+  MountingCoordinator::Shared findMountingCoordinator(
+      SurfaceId surfaceId) const;
+
 #pragma mark - Delegate
 
   /*