respond to review

onflow · Oct 17, 2023 · 5defa03 · 5defa03
1 parent 8363693
commit 5defa03
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 15 deletions.
diff --git a/runtime/sema/check_function.go b/runtime/sema/check_function.go
@@ -187,20 +187,23 @@ func (checker *Checker) checkFunction(
 			functionActivation.InitializationInfo = initializationInfo
 
 			if functionBlock != nil {
-				oldMappedAccess := checker.entitlementMappingInScope
-				if mappedAccess, isMappedAccess := access.(*EntitlementMapAccess); isMappedAccess {
-					checker.entitlementMappingInScope = mappedAccess.Type
-				}
-
-				checker.InNewPurityScope(functionType.Purity == FunctionPurityView, func() {
-					checker.visitFunctionBlock(
-						functionBlock,
-						functionType.ReturnTypeAnnotation,
-						checkResourceLoss,
-					)
-				})
-
-				checker.entitlementMappingInScope = oldMappedAccess
+				func() {
+					oldMappedAccess := checker.entitlementMappingInScope
+					if mappedAccess, isMappedAccess := access.(*EntitlementMapAccess); isMappedAccess {
+						checker.entitlementMappingInScope = mappedAccess.Type
+					} else {
+						checker.entitlementMappingInScope = nil
+					}
+					defer func() { checker.entitlementMappingInScope = oldMappedAccess }()
+
+					checker.InNewPurityScope(functionType.Purity == FunctionPurityView, func() {
+						checker.visitFunctionBlock(
+							functionBlock,
+							functionType.ReturnTypeAnnotation,
+							checkResourceLoss,
+						)
+					})
+				}()
 
 				if mustExit {
 					returnType := functionType.ReturnTypeAnnotation.Type

diff --git a/runtime/sema/checker.go b/runtime/sema/checker.go
@@ -1263,6 +1263,7 @@ func (checker *Checker) functionType(
 	} else {
 		checker.entitlementMappingInScope = nil
 	}
+	defer func() { checker.entitlementMappingInScope = oldMappedAccess }()
 
 	convertedParameters := checker.parameters(parameterList)
 
@@ -1271,7 +1272,6 @@ func (checker *Checker) functionType(
 		convertedReturnTypeAnnotation =
 			checker.ConvertTypeAnnotation(returnTypeAnnotation)
 	}
-	checker.entitlementMappingInScope = oldMappedAccess
 
 	return &FunctionType{
 		Purity:               PurityFromAnnotation(purity),