[FVM] Recover NFT contract

cmd/util/ledger/migrations/contract_checking_migration_test.go

@@ -19,12 +19,12 @@ import (
 	"github.com/onflow/flow-go/model/flow"
 )
 
-func oldExampleTokenCode(fungibleTokenAddress flow.Address) string {
+func oldExampleFungibleTokenCode(fungibleTokenAddress flow.Address) string {
 	return fmt.Sprintf(
 		`
           import FungibleToken from 0x%s
 
-          pub contract ExampleToken: FungibleToken {
+          pub contract ExampleFungibleToken: FungibleToken {
              pub var totalSupply: UFix64
 
              pub resource Vault: FungibleToken.Provider, FungibleToken.Receiver, FungibleToken.Balance {
@@ -68,6 +68,40 @@ func oldExampleTokenCode(fungibleTokenAddress flow.Address) string {
 	)
 }
 
+func oldExampleNonFungibleTokenCode(fungibleTokenAddress flow.Address) string {
+	return fmt.Sprintf(
+		`
+          import NonFungibleToken from 0x%s
+
+          pub contract ExampleNFT: NonFungibleToken {
+
+              /// Total supply of ExampleNFTs in existence
+              pub var totalSupply: UInt64
+
+              /// The core resource that represents a Non Fungible Token.
+              /// New instances will be created using the NFTMinter resource
+              /// and stored in the Collection resource
+              ///
+              pub resource NFT: NonFungibleToken.INFT {
+
+                  /// The unique ID that each NFT has
+                  pub let id: UInt64
+
+                  init(id: UInt64) {
+                      self.id = id
+                  }
+              }
+
+              init() {
+                  // Initialize the total supply
+                  self.totalSupply = 0
+              }
+          }
+        `,
+		fungibleTokenAddress.Hex(),
+	)
+}
+
 func TestContractCheckingMigrationProgramRecovery(t *testing.T) {
 
 	t.Parallel()
@@ -114,15 +148,27 @@ func TestContractCheckingMigrationProgramRecovery(t *testing.T) {
 		systemContracts.FungibleToken,
 		coreContracts.FungibleToken(env),
 	)
+	addSystemContract(
+		systemContracts.NonFungibleToken,
+		coreContracts.NonFungibleToken(env),
+	)
 
-	// Use an old version of the ExampleToken contract,
+	// Use an old version of the ExampleFungibleToken contract,
+	// and "deploy" it at some arbitrary, high (i.e. non-system) address
+	exampleAddress, err := chainID.Chain().AddressAtIndex(1000)
+	require.NoError(t, err)
+	addContract(
+		exampleAddress,
+		"ExampleFungibleToken",
+		[]byte(oldExampleFungibleTokenCode(systemContracts.FungibleToken.Address)),
+	)
+	// Use an old version of the ExampleNonFungibleToken contract,
 	// and "deploy" it at some arbitrary, high (i.e. non-system) address
-	exampleTokenAddress, err := chainID.Chain().AddressAtIndex(1000)
 	require.NoError(t, err)
 	addContract(
-		exampleTokenAddress,
-		"ExampleToken",
-		[]byte(oldExampleTokenCode(systemContracts.FungibleToken.Address)),
+		exampleAddress,
+		"ExampleNonFungibleToken",
+		[]byte(oldExampleNonFungibleTokenCode(systemContracts.NonFungibleToken.Address)),
 	)
 
 	for address, addressContracts := range contracts {
@@ -177,6 +223,11 @@ func TestContractCheckingMigrationProgramRecovery(t *testing.T) {
 
 	assert.Equal(t,
 		[]any{
+			contractCheckingSuccess{
+				AccountAddress: common.Address(systemContracts.NonFungibleToken.Address),
+				ContractName:   systemcontracts.ContractNameNonFungibleToken,
+				Code:           string(coreContracts.NonFungibleToken(env)),
+			},
 			contractCheckingSuccess{
 				AccountAddress: common.Address(systemContracts.ViewResolver.Address),
 				ContractName:   systemcontracts.ContractNameViewResolver,
@@ -193,9 +244,14 @@ func TestContractCheckingMigrationProgramRecovery(t *testing.T) {
 				Code:           string(coreContracts.FungibleToken(env)),
 			},
 			contractCheckingSuccess{
-				AccountAddress: common.Address(exampleTokenAddress),
-				ContractName:   "ExampleToken",
-				Code:           oldExampleTokenCode(systemContracts.FungibleToken.Address),
+				AccountAddress: common.Address(exampleAddress),
+				ContractName:   "ExampleFungibleToken",
+				Code:           oldExampleFungibleTokenCode(systemContracts.FungibleToken.Address),
+			},
+			contractCheckingSuccess{
+				AccountAddress: common.Address(exampleAddress),
+				ContractName:   "ExampleNonFungibleToken",
+				Code:           oldExampleNonFungibleTokenCode(systemContracts.NonFungibleToken.Address),
 			},
 		},
 		reporter.entries,

fvm/environment/program_recovery.go

@@ -29,16 +29,23 @@ func RecoverProgram(
 	sc := systemcontracts.SystemContractsForChain(chainID)
 
 	fungibleTokenAddress := common.Address(sc.FungibleToken.Address)
+	nonFungibleTokenAddress := common.Address(sc.NonFungibleToken.Address)
 
-	if !isFungibleTokenContract(program, fungibleTokenAddress) {
-		return nil, nil
-	}
+	var code string
 
-	contractName := addressLocation.Name
+	switch {
+	case isFungibleTokenContract(program, fungibleTokenAddress):
+		code = RecoveredFungibleTokenCode(fungibleTokenAddress, addressLocation.Name)
 
-	code := RecoveredFungibleTokenCode(fungibleTokenAddress, contractName)
+	case isNonFungibleTokenContract(program, nonFungibleTokenAddress):
+		code = RecoveredNonFungibleTokenCode(nonFungibleTokenAddress, addressLocation.Name)
+	}
+
+	if code != "" {
+		return parser.ParseProgram(memoryGauge, []byte(code), parser.Config{})
+	}
 
-	return parser.ParseProgram(memoryGauge, []byte(code), parser.Config{})
+	return nil, nil
 }
 
 func RecoveredFungibleTokenCode(fungibleTokenAddress common.Address, contractName string) string {
@@ -59,12 +66,12 @@ func RecoveredFungibleTokenCode(fungibleTokenAddress common.Address, contractNam
 
               access(all)
               view fun getContractViews(resourceType: Type?): [Type] {
-                  panic("getContractViews is not implemented")
+                  panic("getContractViews is not available in recovered program")
               }
 
               access(all)
               fun resolveContractView(resourceType: Type?, viewType: Type): AnyStruct? {
-                  panic("resolveContractView is not implemented")
+                  panic("resolveContractView is not available in recovered program")
               }
 
               access(all)
@@ -79,38 +86,38 @@ func RecoveredFungibleTokenCode(fungibleTokenAddress common.Address, contractNam
 
                   access(FungibleToken.Withdraw)
                   fun withdraw(amount: UFix64): @{FungibleToken.Vault} {
-                      panic("withdraw is not implemented")
+                      panic("withdraw is not available in recovered program")
                   }
 
                   access(all)
                   view fun isAvailableToWithdraw(amount: UFix64): Bool {
-                      panic("isAvailableToWithdraw is not implemented")
+                      panic("isAvailableToWithdraw is not available in recovered program")
                   }
 
                   access(all)
                   fun deposit(from: @{FungibleToken.Vault}) {
-                      panic("deposit is not implemented")
+                      panic("deposit is not available in recovered program")
                   }
 
                   access(all)
                   fun createEmptyVault(): @{FungibleToken.Vault} {
-                      panic("createEmptyVault is not implemented")
+                      panic("createEmptyVault is not available in recovered program")
                   }
 
                   access(all)
                   view fun getViews(): [Type] {
-                      panic("getViews is not implemented")
+                      panic("getViews is not available in recovered program")
                   }
 
                   access(all)
                   fun resolveView(_ view: Type): AnyStruct? {
-                      panic("resolveView is not implemented")
+                      panic("resolveView is not available in recovered program")
                   }
               }
 
               access(all)
               fun createEmptyVault(vaultType: Type): @{FungibleToken.Vault} {
-                  panic("createEmptyVault is not implemented")
+                  panic("createEmptyVault is not available in recovered program")
               }
           }
         `,
@@ -119,6 +126,62 @@ func RecoveredFungibleTokenCode(fungibleTokenAddress common.Address, contractNam
 	)
 }
 
+func RecoveredNonFungibleTokenCode(nonFungibleTokenAddress common.Address, contractName string) string {
+	return fmt.Sprintf(
+		//language=Cadence
+		`
+          import NonFungibleToken from %s
+
+          access(all)
+          contract %s: NonFungibleToken {
+
+              access(all)
+              resource NFT: NonFungibleToken.NFT {
+
+                  access(all)
+                  let id: UInt64
+
+                  init(id: UInt64) {
+                      self.id = id
+                  }
+
+                  access(all)
+                  view fun getViews(): [Type] {
+                      panic("getViews is not available in recovered program")
+                  }
+
+                  access(all)
+                  fun resolveView(_ view: Type): AnyStruct? {
+                      panic("resolveView is not available in recovered program")
+                  }
+
+                  access(all)
+                  fun createEmptyCollection(): @{NonFungibleToken.Collection} {
+                      panic("createEmptyCollection is not available in recovered program")
+                  }
+              }
+
+              access(all)
+              view fun getContractViews(resourceType: Type?): [Type] {
+                  panic("getContractViews is not available in recovered program")
+              }
+
+              access(all)
+              fun resolveContractView(resourceType: Type?, viewType: Type): AnyStruct? {
+                  panic("resolveContractView is not available in recovered program")
+              }
+
+              access(all)
+              fun createEmptyCollection(nftType: Type): @{NonFungibleToken.Collection} {
+                  panic("createEmptyCollection is not available in recovered program")
+              }
+          }
+        `,
+		nonFungibleTokenAddress.HexWithPrefix(),
+		contractName,
+	)
+}
+
 func importsAddressLocation(program *ast.Program, address common.Address, name string) bool {
 	importDeclarations := program.ImportDeclarations()
 
@@ -169,6 +232,11 @@ const fungibleTokenTypeTotalSupplyFieldName = "totalSupply"
 const fungibleTokenVaultTypeIdentifier = "Vault"
 const fungibleTokenVaultTypeBalanceFieldName = "balance"
 
+const nonFungibleTokenTypeIdentifier = "NonFungibleToken"
+const nonFungibleTokenTypeTotalSupplyFieldName = "totalSupply"
+const nonFungibleTokenNFTTypeIdentifier = "NFT"
+const nonFungibleTokenNFTTypeIDFieldName = "id"
+
 func isFungibleTokenContract(program *ast.Program, fungibleTokenAddress common.Address) bool {
 
 	// Check if the contract imports the FungibleToken contract
@@ -218,6 +286,54 @@ func isFungibleTokenContract(program *ast.Program, fungibleTokenAddress common.A
 	return true
 }
 
+func isNonFungibleTokenContract(program *ast.Program, nonFungibleTokenAddress common.Address) bool {
+
+	// Check if the contract imports the NonFungibleToken contract
+	if !importsAddressLocation(program, nonFungibleTokenAddress, nonFungibleTokenTypeIdentifier) {
+		return false
+	}
+
+	contractDeclaration := program.SoleContractDeclaration()
+	if contractDeclaration == nil {
+		return false
+	}
+
+	// Check if the contract implements the NonFungibleToken interface
+	if !declaresConformanceTo(contractDeclaration, nonFungibleTokenTypeIdentifier) {
+		return false
+	}
+
+	// Check if the contract has a totalSupply field
+	totalSupplyFieldDeclaration := getField(contractDeclaration, nonFungibleTokenTypeTotalSupplyFieldName)
+	if totalSupplyFieldDeclaration == nil {
+		return false
+	}
+
+	// Check if the totalSupply field is of type UInt64
+	if !isNominalType(totalSupplyFieldDeclaration.TypeAnnotation.Type, sema.UInt64TypeName) {
+		return false
+	}
+
+	// Check if the contract has an NFT resource
+	nftDeclaration := contractDeclaration.Members.CompositesByIdentifier()[nonFungibleTokenNFTTypeIdentifier]
+	if nftDeclaration == nil {
+		return false
+	}
+
+	// Check if the NFT resource has an id field
+	idFieldDeclaration := getField(nftDeclaration, nonFungibleTokenNFTTypeIDFieldName)
+	if idFieldDeclaration == nil {
+		return false
+	}
+
+	// Check if the id field is of type UInt64
+	if !isNominalType(idFieldDeclaration.TypeAnnotation.Type, sema.UInt64TypeName) {
+		return false
+	}
+
+	return true
+}
+
 func getField(declaration *ast.CompositeDeclaration, name string) *ast.FieldDeclaration {
 	for _, fieldDeclaration := range declaration.Members.Fields() {
 		if fieldDeclaration.Identifier.Identifier == name {