diff --git a/Source/ASDisplayNode+Layout.mm b/Source/ASDisplayNode+Layout.mm
index e446cf6b6..f34628976 100644
--- a/Source/ASDisplayNode+Layout.mm
+++ b/Source/ASDisplayNode+Layout.mm
@@ -317,10 +317,10 @@ - (void)_locked_measureNodeWithBoundsIfNecessary:(CGRect)bounds
   [self cancelLayoutTransition];
   
   BOOL didCreateNewContext = NO;
-  ASLayoutElementContext context = ASLayoutElementGetCurrentContext();
-  if (ASLayoutElementContextIsNull(context)) {
-    context = ASLayoutElementContextMake(ASLayoutElementContextDefaultTransitionID);
-    ASLayoutElementSetCurrentContext(context);
+  ASLayoutElementContext *context = ASLayoutElementGetCurrentContext();
+  if (context == nil) {
+    context = [[ASLayoutElementContext alloc] init];
+    ASLayoutElementPushContext(context);
     didCreateNewContext = YES;
   }
   
@@ -341,7 +341,7 @@ - (void)_locked_measureNodeWithBoundsIfNecessary:(CGRect)bounds
   }
   
   if (didCreateNewContext) {
-    ASLayoutElementClearCurrentContext();
+    ASLayoutElementPopContext();
   }
   
   // If our new layout's desired size for self doesn't match current size, ask our parent to update it.
@@ -457,8 +457,8 @@ - (BOOL)_isLayoutTransitionInvalid
 {
   ASDN::MutexLocker l(__instanceLock__);
   if (ASHierarchyStateIncludesLayoutPending(_hierarchyState)) {
-    ASLayoutElementContext context = ASLayoutElementGetCurrentContext();
-    if (ASLayoutElementContextIsNull(context) || _pendingTransitionID != context.transitionID) {
+    ASLayoutElementContext *context = ASLayoutElementGetCurrentContext();
+    if (context == nil || _pendingTransitionID != context.transitionID) {
       return YES;
     }
   }
@@ -546,8 +546,10 @@ - (void)transitionLayoutWithSizeRange:(ASSizeRange)constrainedSize
     ASLayout *newLayout;
     {
       ASDN::MutexLocker l(__instanceLock__);
-      
-      ASLayoutElementSetCurrentContext(ASLayoutElementContextMake(transitionID));
+
+      ASLayoutElementContext *ctx = [[ASLayoutElementContext alloc] init];
+      ctx.transitionID = transitionID;
+      ASLayoutElementPushContext(ctx);
 
       BOOL automaticallyManagesSubnodesDisabled = (self.automaticallyManagesSubnodes == NO);
       self.automaticallyManagesSubnodes = YES; // Temporary flag for 1.9.x
@@ -558,7 +560,7 @@ - (void)transitionLayoutWithSizeRange:(ASSizeRange)constrainedSize
         self.automaticallyManagesSubnodes = NO; // Temporary flag for 1.9.x
       }
       
-      ASLayoutElementClearCurrentContext();
+      ASLayoutElementPopContext();
     }
     
     if (isCancelled()) {
diff --git a/Source/Layout/ASLayoutElement.mm b/Source/Layout/ASLayoutElement.mm
index d31caaae9..6909821a6 100644
--- a/Source/Layout/ASLayoutElement.mm
+++ b/Source/Layout/ASLayoutElement.mm
@@ -30,69 +30,62 @@
 
 #pragma mark - ASLayoutElementContext
 
-CGFloat const ASLayoutElementParentDimensionUndefined = NAN;
-CGSize const ASLayoutElementParentSizeUndefined = {ASLayoutElementParentDimensionUndefined, ASLayoutElementParentDimensionUndefined};
-
-int32_t const ASLayoutElementContextInvalidTransitionID = 0;
-int32_t const ASLayoutElementContextDefaultTransitionID = ASLayoutElementContextInvalidTransitionID + 1;
-
-static inline ASLayoutElementContext _ASLayoutElementContextMake(int32_t transitionID)
-{
-  struct ASLayoutElementContext context;
-  context.transitionID = transitionID;
-  return context;
-}
+@implementation ASLayoutElementContext
 
-static inline BOOL _IsValidTransitionID(int32_t transitionID)
+- (instancetype)init
 {
-  return transitionID > ASLayoutElementContextInvalidTransitionID;
+  if (self = [super init]) {
+    _transitionID = ASLayoutElementContextDefaultTransitionID;
+  }
+  return self;
 }
 
-struct ASLayoutElementContext const ASLayoutElementContextNull = _ASLayoutElementContextMake(ASLayoutElementContextInvalidTransitionID);
+@end
 
-BOOL ASLayoutElementContextIsNull(struct ASLayoutElementContext context)
-{
-  return !_IsValidTransitionID(context.transitionID);
-}
+CGFloat const ASLayoutElementParentDimensionUndefined = NAN;
+CGSize const ASLayoutElementParentSizeUndefined = {ASLayoutElementParentDimensionUndefined, ASLayoutElementParentDimensionUndefined};
 
-ASLayoutElementContext ASLayoutElementContextMake(int32_t transitionID)
-{
-  NSCAssert(_IsValidTransitionID(transitionID), @"Invalid transition ID");
-  return _ASLayoutElementContextMake(transitionID);
-}
+int32_t const ASLayoutElementContextInvalidTransitionID = 0;
+int32_t const ASLayoutElementContextDefaultTransitionID = ASLayoutElementContextInvalidTransitionID + 1;
 
 pthread_key_t ASLayoutElementContextKey;
 
+static void ASLayoutElementDestructor(void *p) {
+  if (p != NULL) {
+    ASDisplayNodeCFailAssert(@"Thread exited without clearing layout element context!");
+    CFBridgingRelease(p);
+  }
+};
+
 // pthread_key_create must be called before the key can be used. This function does that.
 void ASLayoutElementContextEnsureKey()
 {
   static dispatch_once_t onceToken;
   dispatch_once(&onceToken, ^{
-    pthread_key_create(&ASLayoutElementContextKey, free);
+    pthread_key_create(&ASLayoutElementContextKey, ASLayoutElementDestructor);
   });
 }
 
-void ASLayoutElementSetCurrentContext(struct ASLayoutElementContext context)
+void ASLayoutElementPushContext(ASLayoutElementContext *context)
 {
   ASLayoutElementContextEnsureKey();
-  ASDisplayNodeCAssert(pthread_getspecific(ASLayoutElementContextKey) == NULL, @"Nested ASLayoutElementContexts aren't supported.");
-  pthread_setspecific(ASLayoutElementContextKey, new ASLayoutElementContext(context));
+  // NOTE: It would be easy to support nested contexts – just use an NSMutableArray here.
+  ASDisplayNodeCAssertNil(ASLayoutElementGetCurrentContext(), @"Nested ASLayoutElementContexts aren't supported.");
+  pthread_setspecific(ASLayoutElementContextKey, CFBridgingRetain(context));
 }
 
-struct ASLayoutElementContext ASLayoutElementGetCurrentContext()
+ASLayoutElementContext *ASLayoutElementGetCurrentContext()
 {
   ASLayoutElementContextEnsureKey();
-  auto heapCtx = (ASLayoutElementContext *)pthread_getspecific(ASLayoutElementContextKey);
-  return (heapCtx ? *heapCtx : ASLayoutElementContextNull);
+  // Don't retain here. Caller will retain if it wants to!
+  return (__bridge __unsafe_unretained ASLayoutElementContext *)pthread_getspecific(ASLayoutElementContextKey);
 }
 
-void ASLayoutElementClearCurrentContext()
+void ASLayoutElementPopContext()
 {
   ASLayoutElementContextEnsureKey();
-  auto heapCtx = (ASLayoutElementContext *)pthread_getspecific(ASLayoutElementContextKey);
-  if (heapCtx != NULL) {
-    delete heapCtx;
-  }
+  ASDisplayNodeCAssertNotNil(ASLayoutElementGetCurrentContext(), @"Attempt to pop context when there wasn't a context!");
+  CFBridgingRelease(pthread_getspecific(ASLayoutElementContextKey));
   pthread_setspecific(ASLayoutElementContextKey, NULL);
 }
 
diff --git a/Source/Layout/ASLayoutElementPrivate.h b/Source/Layout/ASLayoutElementPrivate.h
index a7f4b16f1..68c46de61 100644
--- a/Source/Layout/ASLayoutElementPrivate.h
+++ b/Source/Layout/ASLayoutElementPrivate.h
@@ -23,26 +23,25 @@
 
 #pragma mark - ASLayoutElementContext
 
-struct ASLayoutElementContext {
-  int32_t transitionID;
-};
+NS_ASSUME_NONNULL_BEGIN
+
+AS_SUBCLASSING_RESTRICTED
+@interface ASLayoutElementContext : NSObject
+@property (nonatomic) int32_t transitionID;
+@end
 
 extern int32_t const ASLayoutElementContextInvalidTransitionID;
 
 extern int32_t const ASLayoutElementContextDefaultTransitionID;
 
-extern struct ASLayoutElementContext const ASLayoutElementContextNull;
-
-extern BOOL ASLayoutElementContextIsNull(struct ASLayoutElementContext context);
-
-extern struct ASLayoutElementContext ASLayoutElementContextMake(int32_t transitionID);
-
-extern void ASLayoutElementSetCurrentContext(struct ASLayoutElementContext context);
+// Does not currently support nesting – there must be no current context.
+extern void ASLayoutElementPushContext(ASLayoutElementContext * context);
 
-extern struct ASLayoutElementContext ASLayoutElementGetCurrentContext();
+extern ASLayoutElementContext * _Nullable ASLayoutElementGetCurrentContext();
 
-extern void ASLayoutElementClearCurrentContext();
+extern void ASLayoutElementPopContext();
 
+NS_ASSUME_NONNULL_END
 
 #pragma mark - ASLayoutElementLayoutDefaults