Repository: onflow/flow-cli (go)

Issue:
Failed to add dependencies by using `flow-c1` from Testnet sources when un-staged contracts exist in the address.
### Problem

Failed to add dependencies by using `flow-c1` from Testnet sources when un-staged contracts exist in the address.

### Steps to Reproduce

`flow-c1 deps add testnet://{AnyAddressWithUnStagedContracts}.{ValidC1.0ContractName}`

### Acceptance Criteria

not to panic and throw error.

### Context

It should be this line that is causing the error, because C1.0 can not read previous Cadence code.
https://github.com/onflow/flow-cli/blob/master/internal/dependencymanager/dependencyinstaller.go#L322-L325


Pull request description:
Fix fetching of dependencies
Fixes #1716

## Description

When using a `flow.json` configuration file with `dependencies` that point to Testnet, then the dependency manager currently implicitly fetches the service account contract `8c5303eaa26202d6.FlowContractAudits`, which has not been updated to Cadence 1.0:

```
❌ Error processing dependency: failed to parse program: Parsing failed:
error: `pub` is no longer a valid access keyword
 --> :1:0
  |
1 | pub contract FlowContractAudits {
  | ^^^
```

This happens if a dependency directly or indirectly depends on a contract in the same account.

This is because currently, when the dependency manager is asked to fetch a dependency (and recursively its dependencies), the manager always fetches all contracts of the account and parses them all, even though only one dependency is requested.

Unless I'm missing something, there is no need to fetch and parse all account contracts, and it is actually incorrect to fail adding a dependency because an unrelated contract in the same account happens to be broken.

Simply fetch the account's contracts, and only parse the requested contract (if it exists).

______

For contributor use:

- [x] Targeted PR against `master` branch
- [x] Linked to Github issue with discussion and accepted design OR link to spec that describes this work
- [x] Code follows the [standards mentioned here](https://github.com/onflow/flow-cli/blob/master/CONTRIBUTING.md#styleguides)
- [x] Updated relevant documentation
- [x] Re-reviewed `Files changed` in the Github PR explorer
- [x] Added appropriate labels


Patch:
diff --git a/go.mod b/go.mod
index 1bcd77e7..6d11cdc6 100644
--- a/go.mod
+++ b/go.mod
@@ -5,7 +5,6 @@ go 1.22
 toolchain go1.22.3
 
 require (
-	github.com/cenkalti/backoff/v4 v4.2.1
 	github.com/charmbracelet/bubbles v0.19.0
 	github.com/charmbracelet/bubbletea v0.27.1
 	github.com/dukex/mixpanel v1.0.1
@@ -15,19 +14,19 @@ require (
 	github.com/gosuri/uilive v0.0.4
 	github.com/logrusorgru/aurora/v4 v4.0.0
 	github.com/manifoldco/promptui v0.9.0
-	github.com/onflow/cadence v1.0.0-preview.50
-	github.com/onflow/cadence-tools/languageserver v1.0.0-preview.38
-	github.com/onflow/cadence-tools/lint v1.0.0-preview.38
-	github.com/onflow/cadence-tools/test v1.0.0-preview.38
+	github.com/onflow/cadence v1.0.0-preview.49
+	github.com/onflow/cadence-tools/languageserver v1.0.0-preview.37
+	github.com/onflow/cadence-tools/lint v1.0.0-preview.37
+	github.com/onflow/cadence-tools/test v1.0.0-preview.37
 	github.com/onflow/contract-updater/lib/go/templates v1.0.1
 	github.com/onflow/fcl-dev-wallet v0.8.0-stable-cadence.1
-	github.com/onflow/flixkit-go v1.2.1-cadence-v1-preview.24
+	github.com/onflow/flixkit-go v1.2.1-cadence-v1-preview.23
 	github.com/onflow/flow-core-contracts/lib/go/templates v1.3.1
-	github.com/onflow/flow-emulator v1.0.0-preview.41
-	github.com/onflow/flow-evm-gateway v0.30.3
-	github.com/onflow/flow-go v0.37.7-0.20240826193109-e211841b59f5
-	github.com/onflow/flow-go-sdk v1.0.0-preview.53
-	github.com/onflow/flowkit/v2 v2.0.0-stable-cadence-alpha.32
+	github.com/onflow/flow-emulator v1.0.0-preview.40
+	github.com/onflow/flow-evm-gateway v0.30.1
+	github.com/onflow/flow-go v0.37.7-0.20240822174309-8b4fff2114d4
+	github.com/onflow/flow-go-sdk v1.0.0-preview.51
+	github.com/onflow/flowkit/v2 v2.0.0-stable-cadence-alpha.31
 	github.com/onflow/go-ethereum v1.14.7
 	github.com/onflowser/flowser/v3 v3.2.1-0.20240131200229-7d4d22715f48
 	github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
@@ -66,6 +65,7 @@ require (
 	github.com/btcsuite/btcd/btcec/v2 v2.3.4 // indirect
 	github.com/btcsuite/btcd/chaincfg/chainhash v1.0.3 // indirect
 	github.com/c-bata/go-prompt v0.2.6 // indirect
+	github.com/cenkalti/backoff/v4 v4.2.1 // indirect
 	github.com/cespare/xxhash v1.1.0 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/charmbracelet/lipgloss v0.13.0 // indirect
diff --git a/go.sum b/go.sum
index 8e10bc4f..1ba65eeb 100644
--- a/go.sum
+++ b/go.sum
@@ -2157,14 +2157,14 @@ github.com/onflow/atree v0.6.1-0.20230711151834-86040b30171f/go.mod h1:xvP61FoOs
 github.com/onflow/atree v0.8.0-rc.6 h1:GWgaylK24b5ta2Hq+TvyOF7X5tZLiLzMMn7lEt59fsA=
 github.com/onflow/atree v0.8.0-rc.6/go.mod h1:yccR+LR7xc1Jdic0mrjocbHvUD7lnVvg8/Ct1AA5zBo=
 github.com/onflow/cadence v1.0.0-M3/go.mod h1:odXGZZ/wGNA5mwT8bC9v8u8EXACHllB2ABSZK65TGL8=
-github.com/onflow/cadence v1.0.0-preview.50 h1:sEfUOG7BXzEqPzB68yZZrG/lkBmHf/o0poYDCY18x3A=
-github.com/onflow/cadence v1.0.0-preview.50/go.mod h1:7wvvecnAZtYOspLOS3Lh+FuAmMeSrXhAWiycC3kQ1UU=
-github.com/onflow/cadence-tools/languageserver v1.0.0-preview.38 h1:/MwSnfAzWZiIyV6CEo8jLnKQRW1vw2hFk0RgO1b1Rh8=
-github.com/onflow/cadence-tools/languageserver v1.0.0-preview.38/go.mod h1:7t7xGMwja5MwiDxYqXUGBzgfmMN3irSespuiYYNFKyM=
-github.com/onflow/cadence-tools/lint v1.0.0-preview.38 h1:ZIYrFkWh2z/iI54W+R9QmdfSN8WG/7tHeBvMx7gOb+0=
-github.com/onflow/cadence-tools/lint v1.0.0-preview.38/go.mod h1:7kEvUs7yVy6y/HPH/8JSfh0w4+peOwPd/a1U7as3l2M=
-github.com/onflow/cadence-tools/test v1.0.0-preview.38 h1:k2QREJ14Seq+Ye7h9V5vqLiZzVMeVuzKP21oMpUyMWw=
-github.com/onflow/cadence-tools/test v1.0.0-preview.38/go.mod h1:90FLNQlrvdIdgiIbYrLr+OYNNcmW3bFmtlKWCXmlbQk=
+github.com/onflow/cadence v1.0.0-preview.49 h1:MXy6qAtFyRSnqSVKpDkh2J1sb2JEryJoYy/5qsaixlw=
+github.com/onflow/cadence v1.0.0-preview.49/go.mod h1:7wvvecnAZtYOspLOS3Lh+FuAmMeSrXhAWiycC3kQ1UU=
+github.com/onflow/cadence-tools/languageserver v1.0.0-preview.37 h1:MA+KM8Hgto4ffUdsHaOqw4iy/gRzALhZdqJr0qO/BN8=
+github.com/onflow/cadence-tools/languageserver v1.0.0-preview.37/go.mod h1:UdlFY+FJk3e94QkDcvvt3gIK/9tr1SOLesRUSJRRSwk=
+github.com/onflow/cadence-tools/lint v1.0.0-preview.37 h1:GffJsraPlJydwTH/N13g4AAOxkkF8fguzhmhxCLpeow=
+github.com/onflow/cadence-tools/lint v1.0.0-preview.37/go.mod h1:6gUUCUyC2ZNx/RSqY+97huXa/Rb857sJvEqCl3ckdmY=
+github.com/onflow/cadence-tools/test v1.0.0-preview.37 h1:TzLg+22QH6x0VETyPTLi0lwaUxHxm9nAMXzU9S89LWY=
+github.com/onflow/cadence-tools/test v1.0.0-preview.37/go.mod h1:IbjA7WlZARTOyXL5AXjVJeYpwzbcMOEY7EW7VU+htKo=
 github.com/onflow/contract-updater/lib/go/templates v1.0.1 h1:xPj898Y8OgLLbXH8+JeKVBV6J+nqPZjiLgGM3Abucto=
 github.com/onflow/contract-updater/lib/go/templates v1.0.1/go.mod h1:OXO6s0X7OW4Q6QTfAfnjoOmibEPgs0psOfMi+tPyzQE=
 github.com/onflow/crypto v0.25.0/go.mod h1:C8FbaX0x8y+FxWjbkHy0Q4EASCDR9bSPWZqlpCLYyVI=
@@ -2172,25 +2172,25 @@ github.com/onflow/crypto v0.25.2 h1:GjHunqVt+vPcdqhxxhAXiMIF3YiLX7gTuTR5O+VG2ns=
 github.com/onflow/crypto v0.25.2/go.mod h1:fY7eLqUdMKV8EGOw301unP8h7PvLVy8/6gVR++/g0BY=
 github.com/onflow/fcl-dev-wallet v0.8.0-stable-cadence.1 h1:IqdUzdqFCSW0klWmA3J9c17ZyQTab9SWcWSLouX6o0Q=
 github.com/onflow/fcl-dev-wallet v0.8.0-stable-cadence.1/go.mod h1:kc42jkiuoPJmxMRFjfbRO9XvnR/3XLheaOerxVMDTiw=
-github.com/onflow/flixkit-go v1.2.1-cadence-v1-preview.24 h1:bBITnm4AcSseozYiFDoFHbI5u/mp3xlccfM/YXHwigc=
-github.com/onflow/flixkit-go v1.2.1-cadence-v1-preview.24/go.mod h1:AndgekHFf9qEqoSwPLkXCgK6hd6lJEV4vOogHiKfXAU=
+github.com/onflow/flixkit-go v1.2.1-cadence-v1-preview.23 h1:Nv9XYhnf3Y0PAFAjSpdejzlpNdjrXMIyztY+HU2EE9A=
+github.com/onflow/flixkit-go v1.2.1-cadence-v1-preview.23/go.mod h1:I0tHJwifIMp/92j9/nXmNypU7HpF2op1XYdN9KXKDcI=
 github.com/onflow/flow-core-contracts/lib/go/contracts v1.3.1 h1:q9tXLIALwQ76bO4bmSrhtTkyc2cZF4/gH11ix9E3F5k=
 github.com/onflow/flow-core-contracts/lib/go/contracts v1.3.1/go.mod h1:u/mkP/B+PbV33tEG3qfkhhBlydSvAKxfLZSfB4lsJHg=
 github.com/onflow/flow-core-contracts/lib/go/templates v1.3.1 h1:FfhMBAb78p6VAWkJ+iqdKLErGQVQgxk5w6DP5ZruWX8=
 github.com/onflow/flow-core-contracts/lib/go/templates v1.3.1/go.mod h1:NgbMOYnMh0GN48VsNKZuiwK7uyk38Wyo8jN9+C9QE30=
-github.com/onflow/flow-emulator v1.0.0-preview.41 h1:Ep8lb4CD5xG7XC55Ne4ZfKHee58osp3jhmMlUYSjZtw=
-github.com/onflow/flow-emulator v1.0.0-preview.41/go.mod h1:selKgeV9J+VC3mwh2tBbv8h1sEDxCldDeRfj1cUqPg4=
-github.com/onflow/flow-evm-gateway v0.30.3 h1:5qKXOBrHHJhDgA1K09gWapFZqbIzw6HetKi9pFUkbvw=
-github.com/onflow/flow-evm-gateway v0.30.3/go.mod h1:EJVJd/Sp2/vxZR6/dyOK4h9Lgll/fqG7GIXK4doWHvA=
+github.com/onflow/flow-emulator v1.0.0-preview.40 h1:VoblJbxqz6xTzgTL1LbbWhuWJ8isONdH8pwQ3Cy1ObM=
+github.com/onflow/flow-emulator v1.0.0-preview.40/go.mod h1:AjbLv8TSsT2hrKF+XudEohbdKNdWomHjH2RlreuxrdU=
+github.com/onflow/flow-evm-gateway v0.30.1 h1:CM7efdGcfdkFvqb2kkWcnpzAWbPrh1zouw7MlsJSvhI=
+github.com/onflow/flow-evm-gateway v0.30.1/go.mod h1:5tn0KTLfiJyok/psERZdNp5bvFIOKgFaXpgRVEFBJO8=
 github.com/onflow/flow-ft/lib/go/contracts v1.0.0 h1:mToacZ5NWqtlWwk/7RgIl/jeKB/Sy/tIXdw90yKHcV0=
 github.com/onflow/flow-ft/lib/go/contracts v1.0.0/go.mod h1:PwsL8fC81cjnUnTfmyL/HOIyHnyaw/JA474Wfj2tl6A=
 github.com/onflow/flow-ft/lib/go/templates v1.0.0 h1:6cMS/lUJJ17HjKBfMO/eh0GGvnpElPgBXx7h5aoWJhs=
 github.com/onflow/flow-ft/lib/go/templates v1.0.0/go.mod h1:uQ8XFqmMK2jxyBSVrmyuwdWjTEb+6zGjRYotfDJ5pAE=
-github.com/onflow/flow-go v0.37.7-0.20240826193109-e211841b59f5 h1:rgX5DJC++uXpeVUiUeOY1fvvQlTWuhrstDz+OTluWIw=
-github.com/onflow/flow-go v0.37.7-0.20240826193109-e211841b59f5/go.mod h1:enp027y6GD0GP/PU9x2Ki5E6X+BSLvNN88Ma4RklZDA=
+github.com/onflow/flow-go v0.37.7-0.20240822174309-8b4fff2114d4 h1:DpC/WAaNDHTbKN/L8HhSwGMGZWPi+C6fUPAYeUjwImc=
+github.com/onflow/flow-go v0.37.7-0.20240822174309-8b4fff2114d4/go.mod h1:iZG3ofIPrPAChSdZxByDGrsBAC5CVDw2hyo4akoQ1NY=
 github.com/onflow/flow-go-sdk v1.0.0-M1/go.mod h1:TDW0MNuCs4SvqYRUzkbRnRmHQL1h4X8wURsCw9P9beo=
-github.com/onflow/flow-go-sdk v1.0.0-preview.53 h1:hF6X5b9gVQoc/wUa5JTcYU6tw1B7aW71XDdfsFhU5Gw=
-github.com/onflow/flow-go-sdk v1.0.0-preview.53/go.mod h1:FtsoOHw5RNmMMwC7jI1hc99ZVJWhnm/gE3OMvFaZjyg=
+github.com/onflow/flow-go-sdk v1.0.0-preview.51 h1:rsvjyEdiot4oKwx2YsEqV2hsXqDrakuVmSUexsU/lSI=
+github.com/onflow/flow-go-sdk v1.0.0-preview.51/go.mod h1:yXinctVJY1VlpQaJhkCUMiC5lvv2kYCCMdMo+FdSV5A=
 github.com/onflow/flow-nft/lib/go/contracts v1.2.1 h1:woAAS5z651sDpi7ihAHll8NvRS9uFXIXkL6xR+bKFZY=
 github.com/onflow/flow-nft/lib/go/contracts v1.2.1/go.mod h1:2gpbza+uzs1k7x31hkpBPlggIRkI53Suo0n2AyA2HcE=
 github.com/onflow/flow-nft/lib/go/templates v1.2.0 h1:JSQyh9rg0RC+D1930BiRXN8lrtMs+ubVMK6aQPon6Yc=
@@ -2198,8 +2198,8 @@ github.com/onflow/flow-nft/lib/go/templates v1.2.0/go.mod h1:p+2hRvtjLUR3MW1NsoJ
 github.com/onflow/flow/protobuf/go/flow v0.3.2-0.20231121210617-52ee94b830c2/go.mod h1:NA2pX2nw8zuaxfKphhKsk00kWLwfd+tv8mS23YXO4Sk=
 github.com/onflow/flow/protobuf/go/flow v0.4.5 h1:6o+pgYGqwXdEhqSJxu2BdnDXkOQVOkfGAb6IiXB+NPM=
 github.com/onflow/flow/protobuf/go/flow v0.4.5/go.mod h1:NA2pX2nw8zuaxfKphhKsk00kWLwfd+tv8mS23YXO4Sk=
-github.com/onflow/flowkit/v2 v2.0.0-stable-cadence-alpha.32 h1:7gzHh7ITIJyBMmqIfXtmu2i+XRFyfsDcdNAheyN2UdQ=
-github.com/onflow/flowkit/v2 v2.0.0-stable-cadence-alpha.32/go.mod h1:0IuMeFtFbv1lCaoF7pB9U2trYjwMzqvuTgSj4mXZX24=
+github.com/onflow/flowkit/v2 v2.0.0-stable-cadence-alpha.31 h1:DzkRlfkwJo1seJp4QlV1Yli13g4rT1kmf2y3WCn8CAg=
+github.com/onflow/flowkit/v2 v2.0.0-stable-cadence-alpha.31/go.mod h1:hd1m2k64ZRlCpGnZ1VFeLRRcCXV186l9+chb3M643h8=
 github.com/onflow/go-ethereum v1.14.7 h1:gg3awYqI02e3AypRdpJKEvNTJ6kz/OhAqRti0h54Wlc=
 github.com/onflow/go-ethereum v1.14.7/go.mod h1:zV14QLrXyYu5ucvcwHUA0r6UaqveqbXaehAVQJlSW+I=
 github.com/onflow/nft-storefront/lib/go/contracts v1.0.0 h1:sxyWLqGm/p4EKT6DUlQESDG1ZNMN9GjPCm1gTq7NGfc=
diff --git a/internal/dependencymanager/dependencyinstaller.go b/internal/dependencymanager/dependencyinstaller.go
index 997d2ae4..5fb562bd 100644
--- a/internal/dependencymanager/dependencyinstaller.go
+++ b/internal/dependencymanager/dependencyinstaller.go
@@ -316,44 +316,40 @@ func (di *DependencyInstaller) fetchDependencies(networkName string, address flo
 		return fmt.Errorf("contracts are nil for account: %s", address)
 	}
 
-	found := false
+	contract, ok := account.Contracts[contractName]
+	if !ok {
+		return fmt.Errorf("contract %s not found at address %s", contractName, address.String())
+	}
 
-	for _, contract := range account.Contracts {
-		program, err := project.NewProgram(contract, nil, "")
-		if err != nil {
-			return fmt.Errorf("failed to parse program: %w", err)
-		}
+	program, err := project.NewProgram(contract, nil, "")
+	if err != nil {
+		return fmt.Errorf("failed to parse program: %w", err)
+	}
 
-		parsedContractName, err := program.Name()
-		if err != nil {
-			return fmt.Errorf("failed to parse contract name: %w", err)
-		}
+	parsedContractName, err := program.Name()
+	if err != nil {
+		return fmt.Errorf("failed to parse contract name: %w", err)
+	}
 
-		if parsedContractName == contractName {
-			found = true
+	if parsedContractName != contractName {
+		return fmt.Errorf("contract name mismatch: expected %s, got %s", contractName, parsedContractName)
+	}
 
-			if err := di.handleFoundContract(networkName, address.String(), assignedName, parsedContractName, program); err != nil {
-				return fmt.Errorf("failed to handle found contract: %w", err)
-			}
+	if err := di.handleFoundContract(networkName, address.String(), assignedName, contractName, program); err != nil {
+		return fmt.Errorf("failed to handle found contract: %w", err)
+	}
 
-			if program.HasAddressImports() {
-				imports := program.AddressImportDeclarations()
-				for _, imp := range imports {
-					contractName := imp.Identifiers[0].String()
-					err := di.fetchDependencies(networkName, flowsdk.HexToAddress(imp.Location.String()), contractName, contractName)
-					if err != nil {
-						return err
-					}
-				}
+	if program.HasAddressImports() {
+		imports := program.AddressImportDeclarations()
+		for _, imp := range imports {
+			contractName := imp.Identifiers[0].String()
+			err := di.fetchDependencies(networkName, flowsdk.HexToAddress(imp.Location.String()), contractName, contractName)
+			if err != nil {
+				return err
 			}
 		}
 	}
 
-	if !found {
-		errMsg := fmt.Sprintf("contract %s not found for account %s on network %s", contractName, address, networkName)
-		di.Logger.Error(errMsg)
-	}
-
 	return nil
 }
 
diff --git a/internal/migrate/staging_validator.go b/internal/migrate/staging_validator.go
index 8e90110e..41ddfada 100644
--- a/internal/migrate/staging_validator.go
+++ b/internal/migrate/staging_validator.go
@@ -25,9 +25,6 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/rs/zerolog"
-	"golang.org/x/exp/slices"
-
 	"github.com/onflow/cadence"
 	"github.com/onflow/cadence/runtime"
 	"github.com/onflow/cadence/runtime/ast"
@@ -38,14 +35,12 @@ import (
 	"github.com/onflow/cadence/runtime/pretty"
 	"github.com/onflow/cadence/runtime/sema"
 	"github.com/onflow/cadence/runtime/stdlib"
-
-	"github.com/onflow/flow-go/cmd/util/ledger/migrations"
-	"github.com/onflow/flow-go/cmd/util/ledger/reporters"
-	"github.com/onflow/flow-go/model/flow"
-
 	"github.com/onflow/contract-updater/lib/go/templates"
 	flowsdk "github.com/onflow/flow-go-sdk"
+	"github.com/onflow/flow-go/cmd/util/ledger/migrations"
+	"github.com/onflow/flow-go/model/flow"
 	"github.com/onflow/flowkit/v2"
+	"golang.org/x/exp/slices"
 
 	"github.com/onflow/flow-cli/internal/util"
 )
@@ -63,13 +58,9 @@ type stagingValidatorImpl struct {
 
 	// Cache for account contract names so we don't have to fetch them multiple times
 	accountContractNames map[common.Address][]string
-
 	// All resolved contract code
 	contracts map[common.Location][]byte
 
-	// Contract codes that are not updated/staged
-	oldCodes map[common.Location][]byte
-
 	// Dependency graph for staged contracts
 	// This root level map holds all nodes
 	graph map[common.Location]node
@@ -181,7 +172,6 @@ func newStagingValidator(flow flowkit.Services) *stagingValidatorImpl {
 		checkingCache:        make(map[common.Location]*cachedCheckingResult),
 		accountContractNames: make(map[common.Address][]string),
 		graph:                make(map[common.Location]node),
-		oldCodes:             make(map[common.Location][]byte),
 	}
 }
 
@@ -196,12 +186,10 @@ func (v *stagingValidatorImpl) Validate(stagedContracts []stagedContractUpdate)
 
 	v.stagedContracts = make(map[common.AddressLocation]stagedContractUpdate)
 	for _, stagedContract := range stagedContracts {
-		stagedContractLocation := stagedContract.DeployLocation
-
-		v.stagedContracts[stagedContractLocation] = stagedContract
+		v.stagedContracts[stagedContract.DeployLocation] = stagedContract
 
 		// Add the contract code to the contracts map for pretty printing
-		v.contracts[stagedContractLocation] = stagedContract.Code
+		v.contracts[stagedContract.SourceLocation] = stagedContract.Code
 	}
 
 	// Load system contracts
@@ -210,72 +198,24 @@ func (v *stagingValidatorImpl) Validate(stagedContracts []stagedContractUpdate)
 	// Parse and check all staged contracts
 	errs := v.checkAllStaged()
 
-	typeRequirements := &migrations.LegacyTypeRequirements{}
-
-	// Extract type requirements from the old codes for all staged contracts.
-	for _, contract := range v.stagedContracts {
-		location := contract.DeployLocation
-
-		// Don't validate contracts with existing errors
-		if errs[location] != nil {
-			continue
-		}
-
-		// Get the account for the contract
-		address := flowsdk.Address(location.Address)
-
-		var account *flowsdk.Account
-		var err error
-
-		err = withRetry(func() error {
-			account, err = v.flow.GetAccount(context.Background(), address)
-			return err
-		})
-		if err != nil {
-			return fmt.Errorf("failed to get account: %w", err)
-		}
-
-		// Get the target contract old code
-		contractName := location.Name
-		oldCode, ok := account.Contracts[contractName]
-		if !ok {
-			return fmt.Errorf("old contract code not found for contract: %s", contractName)
-		}
-		v.oldCodes[location] = oldCode
-
-		migrations.ExtractTypeRequirements(
-			migrations.AddressContract{
-				Location: location,
-				Code:     oldCode,
-			},
-			zerolog.Nop(),
-			reporters.ReportNilWriter{},
-			typeRequirements,
-		)
-	}
-
 	// Validate all contract updates
 	for _, contract := range v.stagedContracts {
-		location := contract.DeployLocation
-
 		// Don't validate contracts with existing errors
-		if errs[location] != nil {
+		if errs[contract.SourceLocation] != nil {
 			continue
 		}
 
 		// Validate the contract update
-		checker := v.checkingCache[location].checker
-		err := v.validateContractUpdate(contract, checker, typeRequirements)
+		checker := v.checkingCache[contract.SourceLocation].checker
+		err := v.validateContractUpdate(contract, checker)
 		if err != nil {
-			errs[location] = err
+			errs[contract.SourceLocation] = err
 		}
 	}
 
 	// Check for any upstream contract update failures
 	for _, contract := range v.stagedContracts {
-		location := contract.DeployLocation
-
-		err := errs[location]
+		err := errs[contract.SourceLocation]
 
 		// We will override any errors other than those related
 		// to missing dependencies, since they are more specific
@@ -293,14 +233,19 @@ func (v *stagingValidatorImpl) Validate(stagedContracts []stagedContractUpdate)
 
 		badDeps := make([]common.Location, 0)
 		v.forEachDependency(contract, func(dependency common.Location) {
-			if errs[dependency] != nil {
+			strLocation, ok := dependency.(common.StringLocation)
+			if !ok {
+				return
+			}
+
+			if errs[strLocation] != nil {
 				badDeps = append(badDeps, dependency)
 			}
 		})
 
 		if len(badDeps) > 0 {
-			errs[location] = &upstreamValidationError{
-				Location:        location,
+			errs[contract.SourceLocation] = &upstreamValidationError{
+				Location:        contract.SourceLocation,
 				BadDependencies: badDeps,
 			}
 		}
@@ -311,7 +256,7 @@ func (v *stagingValidatorImpl) Validate(stagedContracts []stagedContractUpdate)
 		// Map errors to address locations
 		errsByAddress := make(map[common.AddressLocation]error)
 		for _, contract := range v.stagedContracts {
-			err := errs[contract.DeployLocation]
+			err := errs[contract.SourceLocation]
 			if err != nil {
 				errsByAddress[contract.DeployLocation] = err
 			}
@@ -321,13 +266,12 @@ func (v *stagingValidatorImpl) Validate(stagedContracts []stagedContractUpdate)
 	return nil
 }
 
-func (v *stagingValidatorImpl) checkAllStaged() map[common.Location]error {
-	errs := make(map[common.Location]error)
+func (v *stagingValidatorImpl) checkAllStaged() map[common.StringLocation]error {
+	errs := make(map[common.StringLocation]error)
 	for _, contract := range v.stagedContracts {
-		location := contract.DeployLocation
-		_, err := v.checkContract(location)
+		_, err := v.checkContract(contract.SourceLocation)
 		if err != nil {
-			errs[location] = err
+			errs[contract.SourceLocation] = err
 		}
 	}
 
@@ -336,8 +280,6 @@ func (v *stagingValidatorImpl) checkAllStaged() map[common.Location]error {
 	// Note: nodes are not visited more than once so cyclic imports are not an issue
 	// They will be reported, however, by the checker, if they do exist
 	for _, contract := range v.stagedContracts {
-		location := contract.DeployLocation
-
 		// Create a set of all dependencies
 		missingDependencies := make([]common.AddressLocation, 0)
 		v.forEachDependency(contract, func(dependency common.Location) {
@@ -351,7 +293,7 @@ func (v *stagingValidatorImpl) checkAllStaged() map[common.Location]error {
 
 		if len(missingDependencies) > 0 {
 			// If an error exists, only overwrite if it is a checking error
-			existingErr, ok := errs[location]
+			existingErr, ok := errs[contract.SourceLocation]
 			if ok {
 				var existingCheckingErr *sema.CheckerError
 				if !errors.As(existingErr, &existingCheckingErr) {
@@ -359,7 +301,7 @@ func (v *stagingValidatorImpl) checkAllStaged() map[common.Location]error {
 				}
 			}
 
-			errs[location] = &missingDependenciesError{
+			errs[contract.SourceLocation] = &missingDependenciesError{
 				MissingContracts: missingDependencies,
 			}
 		}
@@ -368,11 +310,7 @@ func (v *stagingValidatorImpl) checkAllStaged() map[common.Location]error {
 	return errs
 }
 
-func (v *stagingValidatorImpl) validateContractUpdate(
-	contract stagedContractUpdate,
-	checker *sema.Checker,
-	typeRequirements *migrations.LegacyTypeRequirements,
-) (err error) {
+func (v *stagingValidatorImpl) validateContractUpdate(contract stagedContractUpdate, checker *sema.Checker) (err error) {
 	// Gracefully recover from panics
 	defer func() {
 		if r := recover(); r != nil {
@@ -380,11 +318,21 @@ func (v *stagingValidatorImpl) validateContractUpdate(
 		}
 	}()
 
-	location := contract.DeployLocation
-	contractName := location.Name
+	// Get the account for the contract
+	address := flowsdk.Address(contract.DeployLocation.Address)
+
+	var account *flowsdk.Account
+	err = withRetry(func() error {
+		account, err = v.flow.GetAccount(context.Background(), address)
+		return err
+	})
+	if err != nil {
+		return fmt.Errorf("failed to get account: %w", err)
+	}
 
 	// Get the target contract old code
-	contractCode, ok := v.oldCodes[location]
+	contractName := contract.DeployLocation.Name
+	contractCode, ok := account.Contracts[contractName]
 	if !ok {
 		return fmt.Errorf("old contract code not found for contract: %s", contractName)
 	}
@@ -400,7 +348,7 @@ func (v *stagingValidatorImpl) validateContractUpdate(
 
 	// Check if contract code is valid according to Cadence V1 Update Checker
 	validator := stdlib.NewCadenceV042ToV1ContractUpdateValidator(
-		location,
+		contract.SourceLocation,
 		contractName,
 		&accountContractNamesProviderImpl{
 			resolverFunc: func(address common.Address) ([]string, error) {
@@ -423,6 +371,9 @@ func (v *stagingValidatorImpl) validateContractUpdate(
 		return fmt.Errorf("unsupported network: %s", v.flow.Network().Name)
 	}
 
+	// TODO: extract type requirements from the old contracts
+	typeRequirements := &migrations.LegacyTypeRequirements{}
+
 	validator.WithUserDefinedTypeChangeChecker(
 		migrations.NewUserDefinedTypeChangeCheckerFunc(chainId, typeRequirements),
 	)
@@ -632,11 +583,19 @@ func (v *stagingValidatorImpl) resolveLocation(
 	for i := range resolvedLocations {
 		identifier := identifiers[i]
 
-		resolvedLocation := common.AddressLocation{
+		var resolvedLocation common.Location
+		resovledAddrLocation := common.AddressLocation{
 			Address: addressLocation.Address,
 			Name:    identifier.Identifier,
 		}
 
+		// If the contract one of our staged contract updates, use the source location
+		if stagedUpdate, ok := v.stagedContracts[resovledAddrLocation]; ok {
+			resolvedLocation = stagedUpdate.SourceLocation
+		} else {
+			resolvedLocation = resovledAddrLocation
+		}
+
 		resolvedLocations[i] = runtime.ResolvedLocation{
 			Location:    resolvedLocation,
 			Identifiers: []runtime.Identifier{identifier},
@@ -806,7 +765,7 @@ func (v *stagingValidatorImpl) forEachDependency(
 			}
 		}
 	}
-	traverse(contract.DeployLocation)
+	traverse(contract.SourceLocation)
 }
 
 // Helper for pretty printing errors


Test patch:
diff --git a/internal/migrate/staging_validator_test.go b/internal/migrate/staging_validator_test.go
index a1e15dc0..f91c57cc 100644
--- a/internal/migrate/staging_validator_test.go
+++ b/internal/migrate/staging_validator_test.go
@@ -496,7 +496,7 @@ func Test_StagingValidator(t *testing.T) {
 		// check that error exists & ensure that the local contract names are used (not the deploy locations)
 		fooErr := validatorErr.errors[simpleAddressLocation("0x01.Foo")]
 		require.ErrorContains(t, fooErr, "mismatched types")
-		require.ErrorContains(t, fooErr, "0000000000000001.Foo")
+		require.ErrorContains(t, fooErr, "Foo.cdc")
 
 		// Bar should have an error related to
 		var upstreamErr *upstreamValidationError
@@ -822,115 +822,4 @@ func Test_StagingValidator(t *testing.T) {
 		err := validator.Validate([]stagedContractUpdate{{location, sourceCodeLocation, []byte(newContract)}})
 		require.NoError(t, err)
 	})
-
-	t.Run("with type requirements", func(t *testing.T) {
-		// setup mocks
-		srv := setupValidatorMocks(t, []mockNetworkAccount{
-			{
-				address: flow.HexToAddress("01"),
-				contracts: map[string][]byte{"Foo": []byte(`
-				pub contract interface Foo {
-					pub let bar: @Bar?
-					pub resource Bar {}
-				}`)},
-			},
-			{
-				address: flow.HexToAddress("02"),
-				contracts: map[string][]byte{"Bar": []byte(`
-				import Foo from 0x01
-				pub contract FooImpl: Foo {
-					pub let bar: @Foo.Bar?
-					pub resource BarImpl {}
-					init() {
-						self.bar <- nil
-					}
-				}`)},
-			},
-		})
-
-		validator := newStagingValidator(srv)
-		err := validator.Validate([]stagedContractUpdate{
-			{
-				DeployLocation: simpleAddressLocation("0x01.Foo"),
-				SourceLocation: common.StringLocation("./Foo.cdc"),
-				Code: []byte(`
-				access(all) contract interface Foo {
-					access(all) let bar: @{Bar}?
-					access(all) resource interface Bar {}
-				}`),
-			},
-			{
-				DeployLocation: simpleAddressLocation("0x02.Bar"),
-				SourceLocation: common.StringLocation("./Bar.cdc"),
-				Code: []byte(`
-				import Foo from 0x01
-				access(all) contract FooImpl: Foo {
-					access(all) let bar: @{Foo.Bar}?
-					access(all) resource BarImpl: Foo.Bar {}
-					init() {
-						self.bar <- nil
-					}
-				}`),
-			},
-		})
-
-		require.NoError(t, err)
-	})
-
-	t.Run("contract update with entitlements", func(t *testing.T) {
-		// setup mocks
-		srv := setupValidatorMocks(t, []mockNetworkAccount{
-			{
-				address: flow.HexToAddress("01"),
-				contracts: map[string][]byte{"Foo": []byte(`
-				pub contract Foo {
-					pub resource Bar {}
-				}`)},
-			},
-			{
-				address: flow.HexToAddress("02"),
-				contracts: map[string][]byte{"Test": []byte(`
-				import Foo from 0x01
-				pub contract Test {
-					pub resource R {
-						pub var bar: auth &Foo.Bar?
-						init() {
-							self.bar = nil
-						}
-					}
-				}`)},
-			},
-		})
-
-		validator := newStagingValidator(srv)
-		err := validator.Validate([]stagedContractUpdate{
-			{
-				DeployLocation: simpleAddressLocation("0x01.Foo"),
-				SourceLocation: common.StringLocation("./Foo.cdc"),
-				Code: []byte(`
-				access(all) contract Foo {
-					access(all) resource Bar {
-						access(E) fun foo(){}
-					}
-					access(all) entitlement E
-				}`),
-			},
-			{
-				DeployLocation: simpleAddressLocation("0x02.Test"),
-				SourceLocation: common.StringLocation("./Test.cdc"),
-				Code: []byte(`
-				import Foo from 0x01
-				access(all) contract Test {
-					access(all) resource R {
-						access(all) var bar: auth(Foo.E) &Foo.Bar?
-						init() {
-							self.bar = nil
-						}
-					}
-				}`),
-			},
-		})
-
-		require.NoError(t, err)
-	})
 }