diff --git a/include/drmdisplaycompositor.h b/include/drmdisplaycompositor.h
index 477f226..cfd8f4a 100644
--- a/include/drmdisplaycompositor.h
+++ b/include/drmdisplaycompositor.h
@@ -41,6 +41,15 @@
 
 namespace android {
 
+enum class FlatteningState {
+  kNone,
+  kNotNeeded,
+  kClientRequested,
+  kClientDone,
+  kSerial,
+  kConcurrent
+};
+
 class DrmDisplayCompositor {
  public:
   DrmDisplayCompositor();
@@ -48,6 +57,11 @@
 
   int Init(ResourceManager *resource_manager, int display);
 
+  template <typename Fn>
+  void SetRefreshCallback(Fn &&refresh_cb) {
+    refresh_display_cb_ = std::forward<Fn>(refresh_cb);
+  }
+
   std::unique_ptr<DrmDisplayComposition> CreateComposition() const;
   std::unique_ptr<DrmDisplayComposition> CreateInitializedComposition() const;
   int ApplyComposition(std::unique_ptr<DrmDisplayComposition> composition);
@@ -62,6 +76,9 @@
     return active_composition_->take_out_fence();
   }
 
+  FlatteningState GetFlatteningState() const;
+  bool ShouldFlattenOnClient() const;
+
   std::tuple<uint32_t, uint32_t, int> GetActiveModeResolution();
 
  private:
@@ -90,7 +107,10 @@
 
   void ApplyFrame(std::unique_ptr<DrmDisplayComposition> composition,
                   int status, bool writeback = false);
+
+  bool IsFlatteningNeeded() const;
   int FlattenActiveComposition();
+  int FlattenOnClient();
   int FlattenSerial(DrmConnector *writeback_conn);
   int FlattenConcurrent(DrmConnector *writeback_conn);
   int FlattenOnDisplay(std::unique_ptr<DrmDisplayComposition> &src,
@@ -126,6 +146,10 @@
   int64_t flatten_countdown_;
   std::unique_ptr<Planner> planner_;
   int writeback_fence_;
+
+  FlatteningState flattening_state_;
+
+  std::function<void(int)> refresh_display_cb_;
 };
 }  // namespace android