diff --git a/docs/diagrams/UndoRedoState0.puml b/docs/diagrams/UndoRedoState0.puml
index 43a45903ac9..80054790710 100644
--- a/docs/diagrams/UndoRedoState0.puml
+++ b/docs/diagrams/UndoRedoState0.puml
@@ -7,10 +7,11 @@ skinparam ClassBackgroundColor #FFFFAA
 title Initial state
 
 package States {
-    class State1 as "<u>ab0:AddressBook</u>"
-    class State2 as "<u>ab1:AddressBook</u>"
-    class State3 as "<u>ab2:AddressBook</u>"
+    class State1 as "<u>state0:State</u>"
+    class State2 as "<u>state1:State</u>"
+    class State3 as "<u>state2:State</u>"
 }
+
 State1 -[hidden]right-> State2
 State2 -[hidden]right-> State3
 hide State2
diff --git a/docs/diagrams/UndoRedoState1.puml b/docs/diagrams/UndoRedoState1.puml
index 5a41e9e1651..5dd6714fcba 100644
--- a/docs/diagrams/UndoRedoState1.puml
+++ b/docs/diagrams/UndoRedoState1.puml
@@ -6,10 +6,10 @@ skinparam ClassBackgroundColor #FFFFAA
 
 title After command "delete 5"
 
-package States <<rectangle>> {
-    class State1 as "<u>ab0:AddressBook</u>"
-    class State2 as "<u>ab1:AddressBook</u>"
-    class State3 as "<u>ab2:AddressBook</u>"
+package States {
+    class State1 as "<u>state0:State</u>"
+    class State2 as "<u>state1:State</u>"
+    class State3 as "<u>state2:State</u>"
 }
 
 State1 -[hidden]right-> State2
diff --git a/docs/diagrams/UndoRedoState2.puml b/docs/diagrams/UndoRedoState2.puml
index ad32fce1b0b..ac9653a0572 100644
--- a/docs/diagrams/UndoRedoState2.puml
+++ b/docs/diagrams/UndoRedoState2.puml
@@ -6,10 +6,10 @@ skinparam ClassBackgroundColor #FFFFAA
 
 title After command "add n/David"
 
-package States <<rectangle>> {
-    class State1 as "<u>ab0:AddressBook</u>"
-    class State2 as "<u>ab1:AddressBook</u>"
-    class State3 as "<u>ab2:AddressBook</u>"
+package States {
+    class State1 as "<u>state0:State</u>"
+    class State2 as "<u>state1:State</u>"
+    class State3 as "<u>state2:State</u>"
 }
 
 State1 -[hidden]right-> State2
diff --git a/docs/diagrams/UndoRedoState3.puml b/docs/diagrams/UndoRedoState3.puml
index 9187a690036..3db71c91cd3 100644
--- a/docs/diagrams/UndoRedoState3.puml
+++ b/docs/diagrams/UndoRedoState3.puml
@@ -6,10 +6,10 @@ skinparam ClassBackgroundColor #FFFFAA
 
 title After command "undo"
 
-package States <<rectangle>> {
-    class State1 as "<u>ab0:AddressBook</u>"
-    class State2 as "<u>ab1:AddressBook</u>"
-    class State3 as "<u>ab2:AddressBook</u>"
+package States {
+    class State1 as "<u>state0:State</u>"
+    class State2 as "<u>state1:State</u>"
+    class State3 as "<u>state2:State</u>"
 }
 
 State1 -[hidden]right-> State2
diff --git a/docs/diagrams/UndoRedoState4.puml b/docs/diagrams/UndoRedoState4.puml
index 2bc631ffcd0..5afc3bfd620 100644
--- a/docs/diagrams/UndoRedoState4.puml
+++ b/docs/diagrams/UndoRedoState4.puml
@@ -6,10 +6,10 @@ skinparam ClassBackgroundColor #FFFFAA
 
 title After command "list"
 
-package States <<rectangle>> {
-    class State1 as "<u>ab0:AddressBook</u>"
-    class State2 as "<u>ab1:AddressBook</u>"
-    class State3 as "<u>ab2:AddressBook</u>"
+package States {
+    class State1 as "<u>state0:State</u>"
+    class State2 as "<u>state1:State</u>"
+    class State3 as "<u>state2:State</u>"
 }
 
 State1 -[hidden]right-> State2
diff --git a/docs/diagrams/UndoRedoState5.puml b/docs/diagrams/UndoRedoState5.puml
index e77b04104aa..9920ccb3039 100644
--- a/docs/diagrams/UndoRedoState5.puml
+++ b/docs/diagrams/UndoRedoState5.puml
@@ -6,10 +6,10 @@ skinparam ClassBackgroundColor #FFFFAA
 
 title After command "clear"
 
-package States <<rectangle>> {
-    class State1 as "<u>ab0:AddressBook</u>"
-    class State2 as "<u>ab1:AddressBook</u>"
-    class State3 as "<u>ab3:AddressBook</u>"
+package States {
+    class State1 as "<u>state0:State</u>"
+    class State2 as "<u>state1:State</u>"
+    class State3 as "<u>state3:State</u>"
 }
 
 State1 -[hidden]right-> State2
@@ -18,5 +18,5 @@ State2 -[hidden]right-> State3
 class Pointer as "Current State" #FFFFFF
 
 Pointer -up-> State3
-note right on link: State ab2 deleted.
+note right on link: State state2 deleted.
 @end
diff --git a/docs/diagrams/UndoSequenceDiagram-Logic.puml b/docs/diagrams/UndoSequenceDiagram-Logic.puml
index e57368c5159..68300d53858 100644
--- a/docs/diagrams/UndoSequenceDiagram-Logic.puml
+++ b/docs/diagrams/UndoSequenceDiagram-Logic.puml
@@ -30,9 +30,19 @@ deactivate AddressBookParser
 LogicManager -> UndoCommand : execute()
 activate UndoCommand
 
-UndoCommand -> Model : undoAddressBook()
+UndoCommand -> Model : rollBackHistory()
 activate Model
+Model --> UndoCommand
+deactivate Model
 
+UndoCommand -> Model : getCurrState()
+activate Model
+Model --> UndoCommand : state
+deactivate Model
+
+
+UndoCommand -> Model : restoreState(state)
+activate Model
 Model --> UndoCommand
 deactivate Model
 
diff --git a/docs/diagrams/UndoSequenceDiagram-Model.puml b/docs/diagrams/UndoSequenceDiagram-Model.puml
index 54d83208cb8..ee7990712b1 100644
--- a/docs/diagrams/UndoSequenceDiagram-Model.puml
+++ b/docs/diagrams/UndoSequenceDiagram-Model.puml
@@ -4,13 +4,13 @@ skinparam ArrowFontStyle plain
 
 box Model MODEL_COLOR_T1
 participant ":Model" as Model MODEL_COLOR
-participant ":VersionedAddressBook" as VersionedAddressBook MODEL_COLOR
+participant ":History" as VersionedAddressBook MODEL_COLOR
 end box
 
-[-> Model : undoAddressBook()
+[-> Model : rollBackHistory()
 activate Model
 
-Model -> VersionedAddressBook : undo()
+Model -> VersionedAddressBook : rollBack()
 activate VersionedAddressBook
 
 VersionedAddressBook -> VersionedAddressBook :resetData(ReadOnlyAddressBook)