From ef3cb03d8b2c73fad6743315d25b03935ed544c2 Mon Sep 17 00:00:00 2001
From: Angie Chiang <angiebird@google.com>
Date: Thu, 24 Jan 2019 11:45:29 -0800
Subject: [PATCH] Add get_mv_dist

Given an mv_mode, get_mv_dist() obtains the mv and uses it
to compute distortion.

Change-Id: I58b8c7137b99c2736d651e678f0cd013dbd94877
---
 vp9/encoder/vp9_encoder.c | 36 ++++++++++++++++++++++++++++--------
 vp9/encoder/vp9_mcomp.c   |  5 -----
 vp9/encoder/vp9_mcomp.h   |  6 +++++-
 3 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/vp9/encoder/vp9_encoder.c b/vp9/encoder/vp9_encoder.c
index c2f6004f9..e4f394487 100644
--- a/vp9/encoder/vp9_encoder.c
+++ b/vp9/encoder/vp9_encoder.c
@@ -6014,9 +6014,10 @@ static void mode_estimation(VP9_COMP *cpi, MACROBLOCK *x, MACROBLOCKD *xd,
 }
 
 #if CONFIG_NON_GREEDY_MV
-void set_block_src_pred_buf(MACROBLOCK *x, GF_PICTURE *gf_picture,
-                            int frame_idx, int rf_idx, int mi_row, int mi_col) {
-  MACROBLOCKD *xd = &x->e_mbd;
+static void get_block_src_pred_buf(MACROBLOCKD *xd, GF_PICTURE *gf_picture,
+                                   int frame_idx, int rf_idx, int mi_row,
+                                   int mi_col, struct buf_2d *src,
+                                   struct buf_2d *pre) {
   const int mb_y_offset =
       mi_row * MI_SIZE * xd->cur_buf->y_stride + mi_col * MI_SIZE;
   YV12_BUFFER_CONFIG *ref_frame = NULL;
@@ -6024,11 +6025,11 @@ void set_block_src_pred_buf(MACROBLOCK *x, GF_PICTURE *gf_picture,
   if (ref_frame_idx != -1) {
     ref_frame = gf_picture[ref_frame_idx].frame;
   }
-  x->plane[0].src.buf = xd->cur_buf->y_buffer + mb_y_offset;
-  x->plane[0].src.stride = xd->cur_buf->y_stride;
-  xd->plane[0].pre[0].buf = ref_frame->y_buffer + mb_y_offset;
-  xd->plane[0].pre[0].stride = ref_frame->y_stride;
-  assert(xd->cur_buf->y_stride == ref_frame->y_stride);
+  src->buf = xd->cur_buf->y_buffer + mb_y_offset;
+  src->stride = xd->cur_buf->y_stride;
+  pre->buf = ref_frame->y_buffer + mb_y_offset;
+  pre->stride = ref_frame->y_stride;
+  assert(src->stride == pre->stride);
 }
 
 #define MV_PRECHECK_SIZE 4
@@ -6122,6 +6123,25 @@ int_mv get_mv_from_mv_mode(int mv_mode, VP9_COMP *cpi, TplDepFrame *tpl_frame,
   return mv;
 }
 
+double get_mv_dist(int mv_mode, VP9_COMP *cpi, MACROBLOCKD *xd,
+                   GF_PICTURE *gf_picture, int frame_idx,
+                   TplDepFrame *tpl_frame, int rf_idx, BLOCK_SIZE bsize,
+                   int mi_row, int mi_col) {
+  MV mv = get_mv_from_mv_mode(mv_mode, cpi, tpl_frame, rf_idx, bsize, mi_row,
+                              mi_col)
+              .as_mv;
+  MV full_mv = get_full_mv(&mv);
+  uint32_t sse;
+  struct buf_2d src;
+  struct buf_2d pre;
+  get_block_src_pred_buf(xd, gf_picture, frame_idx, rf_idx, mi_row, mi_col,
+                         &src, &pre);
+  // TODO(angiebird): Consider subpixel when computing the sse.
+  cpi->fn_ptr[bsize].vf(src.buf, src.stride, get_buf_from_mv(&pre, &full_mv),
+                        pre.stride, &sse);
+  return (double)sse;
+}
+
 static double get_feature_score(uint8_t *buf, ptrdiff_t stride, int rows,
                                 int cols) {
   double IxIx = 0;
diff --git a/vp9/encoder/vp9_mcomp.c b/vp9/encoder/vp9_mcomp.c
index 534b15acc..602cc5798 100644
--- a/vp9/encoder/vp9_mcomp.c
+++ b/vp9/encoder/vp9_mcomp.c
@@ -29,11 +29,6 @@
 
 // #define NEW_DIAMOND_SEARCH
 
-static INLINE const uint8_t *get_buf_from_mv(const struct buf_2d *buf,
-                                             const MV *mv) {
-  return &buf->buf[mv->row * buf->stride + mv->col];
-}
-
 void vp9_set_mv_search_range(MvLimits *mv_limits, const MV *mv) {
   int col_min = (mv->col >> 3) - MAX_FULL_PEL_VAL + (mv->col & 7 ? 1 : 0);
   int row_min = (mv->row >> 3) - MAX_FULL_PEL_VAL + (mv->row & 7 ? 1 : 0);
diff --git a/vp9/encoder/vp9_mcomp.h b/vp9/encoder/vp9_mcomp.h
index 6bef88747..779e8d8e7 100644
--- a/vp9/encoder/vp9_mcomp.h
+++ b/vp9/encoder/vp9_mcomp.h
@@ -38,6 +38,11 @@ typedef struct search_site_config {
   int total_steps;
 } search_site_config;
 
+static INLINE const uint8_t *get_buf_from_mv(const struct buf_2d *buf,
+                                             const MV *mv) {
+  return &buf->buf[mv->row * buf->stride + mv->col];
+}
+
 void vp9_init_dsmotion_compensation(search_site_config *cfg, int stride);
 void vp9_init3smotion_compensation(search_site_config *cfg, int stride);
 
@@ -143,7 +148,6 @@ static INLINE MV get_full_mv(const MV *mv) {
   out_mv.col = mv->col >> 3;
   return out_mv;
 }
-
 struct TplDepFrame;
 void vp9_prepare_nb_full_mvs(const struct TplDepFrame *tpl_frame, int mi_row,
                              int mi_col, int rf_idx, BLOCK_SIZE bsize,
-- 
2.40.0