Refactor(subgraph): remove code duplication when loading entities

subgraph/src/DisputeKitClassic.ts

@@ -13,22 +13,22 @@ import { ClassicEvidence } from "../generated/schema";
 import { ensureClassicContributionFromEvent } from "./entities/ClassicContribution";
 import {
   createClassicDisputeFromEvent,
-  loadClassicDisputeWithLog,
+  loadClassicDisputeWithLogs,
 } from "./entities/ClassicDispute";
 import { ensureClassicEvidenceGroup } from "./entities/ClassicEvidenceGroup";
 import {
   createClassicRound,
-  loadClassicRoundWithLog,
+  loadClassicRoundWithLogs,
   updateChoiceFundingFromContributionEvent,
 } from "./entities/ClassicRound";
 import { createClassicVote } from "./entities/ClassicVote";
-import { loadDisputeWithLog } from "./entities/Dispute";
+import { loadDisputeWithLogs } from "./entities/Dispute";
 import { ONE } from "./utils";
 
 export const DISPUTEKIT_ID = "1";
 
 export function handleDisputeCreation(event: DisputeCreation): void {
-  const dispute = loadDisputeWithLog(event.params._coreDisputeID.toString());
+  const dispute = loadDisputeWithLogs(event.params._coreDisputeID.toString());
   if (!dispute) return;
   createClassicDisputeFromEvent(event);
   createClassicRound(dispute.id, dispute.currentRoundIndex);
@@ -52,7 +52,7 @@ export function handleEvidenceEvent(event: EvidenceEvent): void {
 export function handleJustificationEvent(event: JustificationEvent): void {
   const coreDisputeID = event.params._coreDisputeID.toString();
   const classicDisputeID = `${DISPUTEKIT_ID}-${coreDisputeID}`;
-  const classicDispute = loadClassicDisputeWithLog(classicDisputeID);
+  const classicDispute = loadClassicDisputeWithLogs(classicDisputeID);
   if (!classicDispute) return;
   const currentLocalRoundID = `${
     classicDispute.id
@@ -70,7 +70,7 @@ export function handleChoiceFunded(event: ChoiceFunded): void {
   const coreRoundIndex = event.params._coreRoundID.toString();
   const roundID = `${DISPUTEKIT_ID}-${coreDisputeID}-${coreRoundIndex}`;
 
-  const localRound = loadClassicRoundWithLog(roundID);
+  const localRound = loadClassicRoundWithLogs(roundID);
   if (!localRound) return;
 
   const currentFeeRewards = localRound.feeRewards;
@@ -85,7 +85,7 @@ export function handleChoiceFunded(event: ChoiceFunded): void {
     const appealCost = klerosCore.appealCost(BigInt.fromString(coreDisputeID));
     localRound.feeRewards = localRound.feeRewards.minus(appealCost);
 
-    const localDispute = loadClassicDisputeWithLog(
+    const localDispute = loadClassicDisputeWithLogs(
       `${DISPUTEKIT_ID}-${coreDisputeID}`
     );
     if (!localDispute) return;

subgraph/src/KlerosCore.ts

@@ -15,13 +15,16 @@ import { ZERO, ONE } from "./utils";
 import {
   createCourtFromEvent,
   getFeeForJuror,
-  loadCourtWithLog,
+  loadCourtWithLogs,
 } from "./entities/Court";
 import {
   createDisputeKitFromEvent,
   filterSupportedDisputeKits,
 } from "./entities/DisputeKit";
-import { createDisputeFromEvent, loadDisputeWithLog } from "./entities/Dispute";
+import {
+  createDisputeFromEvent,
+  loadDisputeWithLogs,
+} from "./entities/Dispute";
 import { createRoundFromRoundInfo } from "./entities/Round";
 import {
   updateCases,
@@ -50,7 +53,7 @@ export function handleCourtCreated(event: CourtCreated): void {
 export function handleCourtModified(event: CourtModified): void {
   const contract = KlerosCore.bind(event.address);
   const courtContractState = contract.courts(event.params._courtID);
-  const court = loadCourtWithLog(event.params._courtID.toString());
+  const court = loadCourtWithLogs(event.params._courtID.toString());
   if (!court) return;
   court.hiddenVotes = courtContractState.value1;
   court.minStake = courtContractState.value2;
@@ -66,7 +69,7 @@ export function handleDisputeKitCreated(event: DisputeKitCreated): void {
 }
 
 export function handleDisputeKitEnabled(event: DisputeKitEnabled): void {
-  const court = loadCourtWithLog(event.params._courtID.toString());
+  const court = loadCourtWithLogs(event.params._courtID.toString());
   if (!court) return;
   const isEnable = event.params._enable;
   const disputeKitID = event.params._disputeKitID.toString();
@@ -81,7 +84,7 @@ export function handleDisputeCreation(event: DisputeCreation): void {
   const disputeID = event.params._disputeID;
   const disputeStorage = contract.disputes(disputeID);
   const courtID = disputeStorage.value0.toString();
-  const court = loadCourtWithLog(courtID);
+  const court = loadCourtWithLogs(courtID);
   if (!court) return;
   court.numberDisputes = court.numberDisputes.plus(ONE);
   court.save();
@@ -93,7 +96,7 @@ export function handleDisputeCreation(event: DisputeCreation): void {
 
 export function handleNewPeriod(event: NewPeriod): void {
   const disputeID = event.params._disputeID.toString();
-  const dispute = loadDisputeWithLog(disputeID);
+  const dispute = loadDisputeWithLogs(disputeID);
   if (!dispute) return;
   dispute.period = getPeriodName(event.params._period);
   dispute.lastPeriodChange = event.block.timestamp;
@@ -103,7 +106,7 @@ export function handleNewPeriod(event: NewPeriod): void {
 export function handleAppealDecision(event: AppealDecision): void {
   const contract = KlerosCore.bind(event.address);
   const disputeID = event.params._disputeID;
-  const dispute = loadDisputeWithLog(disputeID.toString());
+  const dispute = loadDisputeWithLogs(disputeID.toString());
   if (!dispute) return;
   const newRoundIndex = dispute.currentRoundIndex.plus(ONE);
   const roundID = `${disputeID}-${newRoundIndex.toString()}`;
@@ -118,7 +121,7 @@ export function handleAppealDecision(event: AppealDecision): void {
 export function handleDraw(event: DrawEvent): void {
   createDrawFromEvent(event);
   const disputeID = event.params._disputeID.toString();
-  const dispute = loadDisputeWithLog(disputeID);
+  const dispute = loadDisputeWithLogs(disputeID);
   if (!dispute) return;
   const contract = KlerosCore.bind(event.address);
   updateJurorStake(
@@ -158,9 +161,9 @@ export function handleTokenAndETHShift(event: TokenAndETHShiftEvent): void {
     updateRedistributedPNK(tokenAmount, event.block.timestamp);
   }
   updatePaidETH(ethAmount, event.block.timestamp);
-  const dispute = loadDisputeWithLog(disputeID);
+  const dispute = loadDisputeWithLogs(disputeID);
   if (!dispute) return;
-  const court = loadCourtWithLog(dispute.court);
+  const court = loadCourtWithLogs(dispute.court);
   if (!court) return;
   updateJurorStake(
     jurorAddress,

subgraph/src/entities/ClassicDispute.ts

@@ -1,14 +1,9 @@
-import { log } from "@graphprotocol/graph-ts";
 import { DisputeCreation } from "../../generated/DisputeKitClassic/DisputeKitClassic";
 import { ClassicDispute } from "../../generated/schema";
+import { loadWithLogs } from "../utils";
 
-export function loadClassicDisputeWithLog(id: string): ClassicDispute | null {
-  const classicDispute = ClassicDispute.load(id);
-  if (!classicDispute) {
-    log.error("ClassicDispute not found with id: {}", [id]);
-    return null;
-  }
-  return classicDispute;
+export function loadClassicDisputeWithLogs(id: string): ClassicDispute | null {
+  return loadWithLogs("ClassicDispute", id) as ClassicDispute;
 }
 
 export function createClassicDisputeFromEvent(event: DisputeCreation): void {

subgraph/src/entities/ClassicRound.ts

@@ -1,15 +1,10 @@
-import { BigInt, log } from "@graphprotocol/graph-ts";
+import { BigInt } from "@graphprotocol/graph-ts";
 import { Contribution } from "../../generated/DisputeKitClassic/DisputeKitClassic";
 import { ClassicRound } from "../../generated/schema";
-import { ZERO } from "../utils";
+import { loadWithLogs, ZERO } from "../utils";
 
-export function loadClassicRoundWithLog(id: string): ClassicRound | null {
-  const classicRound = ClassicRound.load(id);
-  if (!classicRound) {
-    log.error("ClassicRound not found with id: {}", [id]);
-    return null;
-  }
-  return classicRound;
+export function loadClassicRoundWithLogs(id: string): ClassicRound | null {
+  return loadWithLogs("ClassicRound", id) as ClassicRound;
 }
 
 export function createClassicRound(
@@ -38,7 +33,7 @@ export function updateChoiceFundingFromContributionEvent(
   const coreRoundIndex = event.params._coreRoundID.toString();
   const roundID = `${disputeKitID}-${coreDisputeID}-${coreRoundIndex}`;
 
-  const classicRound = loadClassicRoundWithLog(roundID);
+  const classicRound = loadClassicRoundWithLogs(roundID);
   if (!classicRound) return;
 
   const choice = event.params._choice;

subgraph/src/entities/Court.ts

@@ -1,17 +1,10 @@
-import { BigInt, log } from "@graphprotocol/graph-ts";
+import { BigInt } from "@graphprotocol/graph-ts";
 import { CourtCreated } from "../../generated/KlerosCore/KlerosCore";
 import { Court } from "../../generated/schema";
-import { ZERO } from "../utils";
+import { loadWithLogs, ZERO } from "../utils";
 
-export function loadCourtWithLog(id: string): Court | null {
-  const court = Court.load(id);
-
-  if (!court) {
-    log.error("Court not found with id: {}", [id]);
-    return null;
-  }
-
-  return court;
+export function loadCourtWithLogs(id: string): Court | null {
+  return loadWithLogs("Court", id) as Court;
 }
 
 export function createCourtFromEvent(event: CourtCreated): void {
@@ -35,7 +28,7 @@ export function createCourtFromEvent(event: CourtCreated): void {
 }
 
 export function getFeeForJuror(id: string): BigInt {
-  const court = loadCourtWithLog(id);
+  const court = loadCourtWithLogs(id);
   if (!court) return ZERO;
   return court.feeForJuror;
 }

subgraph/src/entities/Dispute.ts

@@ -1,20 +1,12 @@
-import { log } from "@graphprotocol/graph-ts";
 import {
   KlerosCore,
   DisputeCreation,
 } from "../../generated/KlerosCore/KlerosCore";
 import { Dispute } from "../../generated/schema";
-import { ZERO } from "../utils";
+import { loadWithLogs, ZERO } from "../utils";
 
-export function loadDisputeWithLog(id: string): Dispute | null {
-  const dispute = Dispute.load(id);
-
-  if (!dispute) {
-    log.error("Dispute not found with id: {}", [id]);
-    return null;
-  }
-
-  return dispute;
+export function loadDisputeWithLogs(id: string): Dispute | null {
+  return loadWithLogs("Dispute", id) as Dispute;
 }
 
 export function createDisputeFromEvent(event: DisputeCreation): void {

subgraph/src/entities/DisputeKit.ts

@@ -1,32 +1,17 @@
-import { Bytes } from "@graphprotocol/graph-ts";
 import { DisputeKitCreated } from "../../generated/KlerosCore/KlerosCore";
 import { DisputeKit } from "../../generated/schema";
-import { ZERO, ONE } from "../utils";
+import { ZERO, ONE, loadWithLogs } from "../utils";
 
-export function ensureDisputeKit(id: string): DisputeKit {
-  let disputeKit = DisputeKit.load(id);
-
-  if (disputeKit) {
-    return disputeKit;
-  }
-  // Should never reach here
-  disputeKit = new DisputeKit(id);
-  disputeKit.parent = "0";
-  disputeKit.address = Bytes.fromHexString("0x0");
-  disputeKit.needsFreezing = false;
-  const parent = ensureDisputeKit("0");
-  disputeKit.depthLevel = id === "0" ? ZERO : parent.depthLevel.plus(ONE);
-  disputeKit.save();
-
-  return disputeKit;
+export function loadDisputeKitWithLogs(id: string): DisputeKit | null {
+  return loadWithLogs("DisputeKit", id) as DisputeKit;
 }
 
 export function createDisputeKitFromEvent(event: DisputeKitCreated): void {
   const disputeKit = new DisputeKit(event.params._disputeKitID.toString());
   disputeKit.parent = event.params._parent.toString();
   disputeKit.address = event.params._disputeKitAddress;
   disputeKit.needsFreezing = false;
-  const parent = DisputeKit.load(event.params._parent.toString());
+  const parent = loadDisputeKitWithLogs(event.params._parent.toString());
   disputeKit.depthLevel = parent ? parent.depthLevel.plus(ONE) : ZERO;
   disputeKit.save();
 }

subgraph/src/entities/JurorTokensPerCourt.ts

@@ -4,7 +4,7 @@ import { JurorTokensPerCourt } from "../../generated/schema";
 import { updateActiveJurors, getDelta } from "../datapoint";
 import { ensureUser } from "./Juror";
 import { ZERO } from "../utils";
-import { loadCourtWithLog } from "./Court";
+import { loadCourtWithLogs } from "./Court";
 
 export function ensureJurorTokensPerCourt(
   jurorAddress: string,
@@ -50,7 +50,7 @@ export function updateJurorStake(
   timestamp: BigInt
 ): void {
   const juror = ensureUser(jurorAddress);
-  const court = loadCourtWithLog(courtID);
+  const court = loadCourtWithLogs(courtID);
   if (!court) return;
   const jurorTokens = ensureJurorTokensPerCourt(jurorAddress, courtID);
   const jurorBalance = contract.getJurorBalance(

subgraph/src/entities/Round.ts

@@ -1,26 +1,10 @@
 import { BigInt } from "@graphprotocol/graph-ts";
 import { KlerosCore__getRoundInfoResult } from "../../generated/KlerosCore/KlerosCore";
 import { Round } from "../../generated/schema";
-import { ZERO } from "../utils";
+import { loadWithLogs } from "../utils";
 
-export function ensureRound(id: string): Round {
-  let round = Round.load(id);
-
-  if (round) {
-    return round;
-  }
-  // Should never reach here
-  round = new Round(id);
-  round.disputeKit = "0";
-  round.tokensAtStakePerJuror = ZERO;
-  round.totalFeesForJurors = ZERO;
-  round.nbVotes = ZERO;
-  round.repartitions = ZERO;
-  round.penalties = ZERO;
-  round.dispute = "0";
-  round.save();
-
-  return round;
+export function loadRoundWithLogs(id: string): Round | null {
+  return loadWithLogs("Round", id) as Round;
 }
 
 export function createRoundFromRoundInfo(
@@ -40,14 +24,3 @@ export function createRoundFromRoundInfo(
   round.dispute = disputeID.toString();
   round.save();
 }
-
-export function filterSupportedDisputeKits(
-  supportedDisputeKits: string[],
-  disputeKitID: string
-): string[] {
-  let result: string[] = [];
-  for (let i = 0; i < supportedDisputeKits.length; i++)
-    if (supportedDisputeKits[i] !== disputeKitID)
-      result = result.concat([supportedDisputeKits[i]]);
-  return result;
-}

subgraph/src/utils.ts

@@ -1,4 +1,15 @@
-import { BigInt } from "@graphprotocol/graph-ts";
+import { BigInt, Entity, store, log } from "@graphprotocol/graph-ts";
 
 export const ZERO = BigInt.fromI32(0);
 export const ONE = BigInt.fromI32(1);
+
+export function loadWithLogs(entityName: string, id: string): Entity | null {
+  const entity = store.get(entityName, id);
+
+  if (!entity) {
+    log.error("{} not found with id: {}", [entityName, id]);
+    return null;
+  }
+
+  return entity;
+}