diff --git a/frontends/p4/simplify.cpp b/frontends/p4/simplify.cpp
index 5b43b5e0189..e372bf0c867 100644
--- a/frontends/p4/simplify.cpp
+++ b/frontends/p4/simplify.cpp
@@ -24,26 +24,20 @@ const IR::Node* DoSimplifyControlFlow::postorder(IR::BlockStatement* statement)
     if (statement->annotations->size() > 0)
         return statement;
     auto parent = getContext()->node;
-    auto statancestor = findContext<IR::Statement>();
-
-    // TODO: either the `parent != nullptr` checks below are redundant or this
-    // if predicate is unsafe.
-
+    CHECK_NULL(parent);
     if (parent->is<IR::SwitchCase>() ||
         parent->is<IR::P4Control>() ||
         parent->is<IR::Function>() ||
         parent->is<IR::P4Action>()) {
-        // Cannot remove block from switch or toplevel control block
+        // Cannot remove these blocks
         return statement;
     }
-    bool withinBlock = statancestor != nullptr && statancestor->is<IR::BlockStatement>();
-
-    // TODO: withinAction will always be false: If parent is a P4Action, the if
-    // statement above would return.
-    bool withinAction = parent != nullptr && parent->is<IR::P4Action>();
+    bool inBlock = findContext<IR::Statement>() != nullptr;
+    bool inState = findContext<IR::ParserState>() != nullptr;
+    if (!(inBlock || inState))
+        return statement;
 
-    bool withinParserState = parent != nullptr && parent->is<IR::ParserState>();
-    if (withinParserState || withinBlock || withinAction) {
+    if (parent->is<IR::BlockStatement>() || parent->is<IR::ParserState>()) {
         // if there are no local declarations we can remove this block
         bool hasDeclarations = false;
         for (auto c : statement->components)
@@ -54,6 +48,7 @@ const IR::Node* DoSimplifyControlFlow::postorder(IR::BlockStatement* statement)
         if (!hasDeclarations)
             return &statement->components;
     }
+
     if (statement->components.empty())
         return new IR::EmptyStatement(statement->srcInfo);
     if (statement->components.size() == 1) {