Refactor property access logic

flintlang · May 7, 2018 · 858a716 · 858a716
1 parent 415de32
commit 858a716
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 28 deletions.
diff --git a/Sources/AST/ASTVisitorContext.swift b/Sources/AST/ASTVisitorContext.swift
@@ -86,4 +86,18 @@ public struct ScopeContext: Equatable {
     let all = localVariables + parameters.map { $0.asVariableDeclaration }
     return all.first(where: { $0.identifier.name == variable })?.type.rawType
   }
+
+  /// Returns the parameter name for the enclosing identifier of the given expression.
+  ///
+  /// For example, when given the expression "a.foo.x", the function will return "a" if "a" is a parameter to the
+  /// function.
+  public func enclosingParameter(expression: Expression, enclosingTypeName: String) -> String? {
+    guard expression.enclosingType != enclosingTypeName,
+      let enclosingIdentifier = expression.enclosingIdentifier,
+      containsParameterDeclaration(for: enclosingIdentifier.name) else {
+      return nil
+    }
+
+    return enclosingIdentifier.name
+  }
 }
diff --git a/Sources/IRGen/Function/IULIAExpression.swift b/Sources/IRGen/Function/IULIAExpression.swift
@@ -96,43 +96,42 @@ struct IULIAPropertyAccess {
   var asLValue: Bool
 
   func rendered(functionContext: FunctionContext) -> String {
-    let lhsOffset: String
-
     let environment = functionContext.environment
     let scopeContext = functionContext.scopeContext
     let enclosingTypeName = functionContext.enclosingTypeName
     let isInStructFunction = functionContext.isInStructFunction
 
     var isMemoryAccess: Bool = false
-
-    if case .identifier(let lhsIdentifier) = lhs {
-      if let enclosingType = lhsIdentifier.enclosingType, let offset = environment.propertyOffset(for: lhsIdentifier.name, enclosingType: enclosingType) {
-        lhsOffset = "\(offset)"
-      } else if functionContext.scopeContext.containsVariableDeclaration(for: lhsIdentifier.name) {
-        lhsOffset = lhsIdentifier.name.mangled
-        isMemoryAccess = true
-      } else {
-        lhsOffset = "\(environment.propertyOffset(for: lhsIdentifier.name, enclosingType: enclosingTypeName)!)"
-      }
-    } else {
-      lhsOffset = IULIAExpression(expression: lhs, asLValue: true).rendered(functionContext: functionContext)
-    }
 
     let lhsType = environment.type(of: lhs, enclosingType: enclosingTypeName, scopeContext: scopeContext)
     let rhsOffset = IULIAPropertyOffset(expression: rhs, enclosingType: lhsType).rendered(functionContext: functionContext)
 
     let offset: String
     if isInStructFunction {
       let enclosingName: String
-      if lhs.enclosingType != functionContext.enclosingTypeName, let enclosingIdentifier = lhs.enclosingIdentifier, functionContext.scopeContext.containsParameterDeclaration(for: enclosingIdentifier.name) {
-        enclosingName = enclosingIdentifier.name
+      if let enclosingParameter = functionContext.scopeContext.enclosingParameter(expression: lhs, enclosingTypeName: functionContext.enclosingTypeName) {
+        enclosingName = enclosingParameter
       } else {
         enclosingName = "flintSelf"
       }
 
       // For struct parameters, access the property by an offset to _flintSelf (the receiver's address).
       offset = IULIARuntimeFunction.addOffset(base: enclosingName.mangled, offset: rhsOffset, inMemory: Mangler.isMem(for: enclosingName).mangled)
     } else {
+      let lhsOffset: String
+      if case .identifier(let lhsIdentifier) = lhs {
+        if let enclosingType = lhsIdentifier.enclosingType, let offset = environment.propertyOffset(for: lhsIdentifier.name, enclosingType: enclosingType) {
+          lhsOffset = "\(offset)"
+        } else if functionContext.scopeContext.containsVariableDeclaration(for: lhsIdentifier.name) {
+          lhsOffset = lhsIdentifier.name.mangled
+          isMemoryAccess = true
+        } else {
+          lhsOffset = "\(environment.propertyOffset(for: lhsIdentifier.name, enclosingType: enclosingTypeName)!)"
+        }
+      } else {
+        lhsOffset = IULIAExpression(expression: lhs, asLValue: true).rendered(functionContext: functionContext)
+      }
+
       offset = IULIARuntimeFunction.addOffset(base: lhsOffset, offset: rhsOffset, inMemory: isMemoryAccess)
     }
 
@@ -181,27 +180,25 @@ struct IULIAAssignment {
     default:
       // LHS refers to a property in storage or memory.
 
-      let isMemoryAccess: Bool
-
-      if let enclosingIdentifier = lhs.enclosingIdentifier, functionContext.scopeContext.containsDeclaration(for: enclosingIdentifier.name) {
-        isMemoryAccess = true
-      } else {
-        isMemoryAccess = false
-      }
-
       let lhsCode = IULIAExpression(expression: lhs, asLValue: true).rendered(functionContext: functionContext)
 
-
       if functionContext.isInStructFunction {
         let enclosingName: String
-        if lhs.enclosingType != functionContext.enclosingTypeName, let enclosingIdentifier = lhs.enclosingIdentifier, functionContext.scopeContext.containsParameterDeclaration(for: enclosingIdentifier.name) {
-          enclosingName = enclosingIdentifier.name
+        if let enclosingParameter = functionContext.scopeContext.enclosingParameter(expression: lhs, enclosingTypeName: functionContext.enclosingTypeName) {
+          enclosingName = enclosingParameter
         } else {
           enclosingName = "flintSelf"
         }
         return IULIARuntimeFunction.store(address: lhsCode, value: rhsCode, inMemory: Mangler.isMem(for: enclosingName).mangled)
       }
 
+      let isMemoryAccess: Bool
+      if let enclosingIdentifier = lhs.enclosingIdentifier, functionContext.scopeContext.containsVariableDeclaration(for: enclosingIdentifier.name) {
+        isMemoryAccess = true
+      } else {
+        isMemoryAccess = false
+      }
+
       return IULIARuntimeFunction.store(address: lhsCode, value: rhsCode, inMemory: isMemoryAccess)
     }
   }