Fix trait filter save issue

diff --git a/MyCollection/Internals/Classes/Piece.lua b/MyCollection/Internals/Classes/Piece.lua
index ee16bc6..b78ef26 100644
--- a/MyCollection/Internals/Classes/Piece.lua
+++ b/MyCollection/Internals/Classes/Piece.lua
@@ -11,14 +11,16 @@ Classes.Piece.__index = Classes.Piece
 local Piece = Classes.Piece

 -- Constructor
-function Piece.New(setId, equipType, traitType, armorType, weaponType)
+function Piece.New(savedTable, setId, equipType, traitType, armorType, weaponType)
     local instance = {}
     setmetatable(instance, Piece)

-    instance:Initialize(setId, equipType, traitType, armorType, weaponType)
+    instance:Initialize(savedTable, setId, equipType, traitType, armorType, weaponType)
     return instance
 end

+Piece.savedTableReference = nil
+
 Piece.references = {}

 -- Properties
@@ -31,6 +33,7 @@ Piece.weaponType = nil
 -- Functions
 function Piece:SetTrait(traitType)
     self.traitType = traitType
+    self.savedTableReference.traitType = traitType
     self:Refresh()
 end

@@ -70,7 +73,9 @@ function Piece:Refresh()
     end
 end

-function Piece:Initialize(setId, equipType, traitType, armorType, weaponType)
+function Piece:Initialize(savedTable, setId, equipType, traitType, armorType, weaponType)
+    self.savedTableReference = savedTable
+
     self.setId = setId
     self.equipType = equipType
     if (traitType == nil) then
@@ -86,16 +91,20 @@ function Piece:Initialize(setId, equipType, traitType, armorType, weaponType)

     self.references = {}

+    if (self.savedTableReference == nil) then
+        self.savedTableReference = {
+             setId = self.setId,
+             equipType = self.equipType,
+             traitType = self.traitType,
+             armorType = self.armorType,
+             weaponType = self.weaponType,
+        }
+     end
+
     self:Refresh()
 end

 -- For saving purpose
-function Piece:CopyForSave()
-    return {
-        setId = self.setId,
-        equipType = self.equipType,
-        traitType = self.traitType,
-        armorType = self.armorType,
-        weaponType = self.weaponType,
-    }
+function Piece:GetSavedTable()
+    return self.savedTableReference
 end
\ No newline at end of file
diff --git a/MyCollection/Internals/Classes/Set.lua b/MyCollection/Internals/Classes/Set.lua
index 4b6167b..140d928 100644
--- a/MyCollection/Internals/Classes/Set.lua
+++ b/MyCollection/Internals/Classes/Set.lua
@@ -128,19 +128,19 @@ function Set:Initialize(savedTable, setId, setName, armorType, orderNumber)
     if savedTable ~= nil and next(savedTable) ~= nil then
         if (next(savedTable.armors) ~= nil) then
             for equipType, item in pairs(savedTable.armors) do
-                self.pieces.armors[equipType] = Classes.Piece.New(item.setId, item.equipType, item.traitType, item.armorType, item.weaponType)
+                self.pieces.armors[equipType] = Classes.Piece.New(savedTable.armors[equipType], item.setId, item.equipType, item.traitType, item.armorType, item.weaponType)
             end
         end

         if (next(savedTable.jewelleries) ~= nil) then
             for equipType, item in pairs(savedTable.jewelleries) do
-                self.pieces.jewelleries[equipType] = Classes.Piece.New(item.setId, item.equipType, item.traitType, item.armorType, item.weaponType)
+                self.pieces.jewelleries[equipType] = Classes.Piece.New(savedTable.armors[equipType], item.setId, item.equipType, item.traitType, item.armorType, item.weaponType)
             end
         end

         if (next(savedTable.weapons) ~= nil) then
             for equipType, item in pairs(savedTable.weapons) do
-                self.pieces.weapons[equipType] = Classes.Piece.New(item.setId, item.equipType, item.traitType, item.armorType, item.weaponType)
+                self.pieces.weapons[equipType] = Classes.Piece.New(savedTable.armors[equipType], item.setId, item.equipType, item.traitType, item.armorType, item.weaponType)
             end
         end
     else
@@ -154,18 +154,18 @@ function Set:Initialize(savedTable, setId, setName, armorType, orderNumber)
         savedTable.weapons = {}

         for _, equipType in pairs(Constants.EquipTypes.Armors) do
-            self.pieces.armors[equipType] = Classes.Piece.New(self.setId, equipType, nil, self.armorType, nil)
-            savedTable.armors[equipType] = self.pieces.armors[equipType]:CopyForSave()
+            self.pieces.armors[equipType] = Classes.Piece.New(nil, self.setId, equipType, nil, self.armorType, nil)
+            savedTable.armors[equipType] = self.pieces.armors[equipType]:GetSavedTable()
         end

         for _, equipType in pairs(Constants.EquipTypes.Jewelleries) do
-            self.pieces.jewelleries[equipType] = Classes.Piece.New(self.setId, equipType, nil, nil, nil)
-            savedTable.jewelleries[equipType] = self.pieces.jewelleries[equipType]:CopyForSave()
+            self.pieces.jewelleries[equipType] = Classes.Piece.New(nil, self.setId, equipType, nil, nil, nil)
+            savedTable.jewelleries[equipType] = self.pieces.jewelleries[equipType]:GetSavedTable()
         end

         for _, weaponType in pairs(Constants.WeaponTypes) do
-            self.pieces.weapons[weaponType] = Classes.Piece.New(self.setId, equipType, nil, nil, weaponType)
-            savedTable.weapons[weaponType] = self.pieces.weapons[weaponType]:CopyForSave()
+            self.pieces.weapons[weaponType] = Classes.Piece.New(nil, self.setId, equipType, nil, nil, weaponType)
+            savedTable.weapons[weaponType] = self.pieces.weapons[weaponType]:GetSavedTable()
         end
     end
 end
\ No newline at end of file
diff --git a/MyCollection/UI/Controls/Grid.lua b/MyCollection/UI/Controls/Grid.lua
index 9f719af..b0f4ef8 100644
--- a/MyCollection/UI/Controls/Grid.lua
+++ b/MyCollection/UI/Controls/Grid.lua
@@ -244,6 +244,7 @@ function Grid:ShowTraitMenu( control )

     AddMenuItem("None", function ()
         control.reference:SetTrait(0)
+        control.traitId = 0
         control.traitName = GetString(traitTextId)
         self:RefreshData()
     end)
@@ -252,6 +253,7 @@ function Grid:ShowTraitMenu( control )
         local traitTextId = _G["MYCOLLECTION_TRAIT_".. control.traitType .. "_" .. key:upper()]
         AddMenuItem(GetString(traitTextId), function ()
             control.reference:SetTrait(trait)
+            control.traitId = trait
             control.traitName = GetString(traitTextId)
             self:RefreshData()
         end)