Fix bugs and clean up

webview/src/experiments/components/table/MergeHeaderGroups.tsx

@@ -13,6 +13,7 @@ export const MergedHeaderGroups: React.FC<{
   onDragUpdate: DragFunction
   onDragStart: DragFunction
   onDragEnd: DragFunction
+  onDrop: DragFunction
   setExpColumnNeedsShadow: (needsShadow: boolean) => void
   root: HTMLElement | null
 }> = ({
@@ -22,6 +23,7 @@ export const MergedHeaderGroups: React.FC<{
   onDragUpdate,
   onDragEnd,
   onDragStart,
+  onDrop,
   root,
   setExpColumnNeedsShadow
 }) => {
@@ -40,7 +42,8 @@ export const MergedHeaderGroups: React.FC<{
           columns={columns}
           onDragEnter={onDragUpdate}
           onDragStart={onDragStart}
-          onDrop={onDragEnd}
+          onDragEnd={onDragEnd}
+          onDrop={onDrop}
           root={root}
         />
       ))}

webview/src/experiments/components/table/Table.test.tsx

@@ -438,6 +438,7 @@ describe('Table', () => {
       )
 
       headers = await getHeaders()
+      console.log('2', Object.assign({}, headers))
 
       expect(headers.indexOf('loss')).toBeGreaterThan(
         headers.indexOf('threshold')

webview/src/experiments/components/table/Table.tsx

@@ -1,4 +1,4 @@
-import React, { useRef, useState } from 'react'
+import React, { useRef, useState, CSSProperties } from 'react'
 import { useSelector } from 'react-redux'
 import cx from 'classnames'
 import styles from './styles.module.scss'
@@ -61,7 +61,10 @@ export const Table: React.FC<InstanceProp> = ({ instance }) => {
   )
 
   return (
-    <div className={styles.tableContainer}>
+    <div
+      className={styles.tableContainer}
+      style={{ '--table-head-height': `${tableHeadHeight}px` } as CSSProperties}
+    >
       <div
         {...getTableProps({
           className: cx(

webview/src/experiments/components/table/TableBody.tsx

@@ -1,4 +1,4 @@
-import React, { CSSProperties } from 'react'
+import React from 'react'
 import cx from 'classnames'
 import { useInView } from 'react-intersection-observer'
 import styles from './styles.module.scss'
@@ -21,9 +21,6 @@ const WorkspaceRowGroupWrapper: React.FC<
 
   return (
     <div
-      style={
-        { '--table-head-height': `${tableHeaderHeight}px` } as CSSProperties
-      }
       ref={ref}
       {...instance.getTableBodyProps({
         className: cx(

webview/src/experiments/components/table/TableHead.tsx

@@ -1,17 +1,19 @@
 import { Experiment } from 'dvc/src/experiments/webview/contract'
 import React, { DragEvent, useRef, useEffect } from 'react'
-import { useSelector } from 'react-redux'
+import { useSelector, useDispatch } from 'react-redux'
 import { HeaderGroup, TableInstance } from 'react-table'
 import { MessageFromWebviewType } from 'dvc/src/webview/contract'
 import styles from './styles.module.scss'
 import { MergedHeaderGroups } from './MergeHeaderGroups'
 import { Indicators } from './Indicators'
+import { setDropTarget } from './headerDropTargetSlice'
 import { useColumnOrder } from '../../hooks/useColumnOrder'
 import { ExperimentsState } from '../../store'
 import { sendMessage } from '../../../shared/vscode'
 import { leafColumnIds, reorderColumnIds } from '../../util/columns'
 import { DragFunction } from '../../../shared/components/dragDrop/Draggable'
 import { getSelectedForPlotsCount } from '../../util/rows'
+
 interface TableHeadProps {
   instance: TableInstance<Experiment>
   root: HTMLElement | null
@@ -34,14 +36,19 @@ export const TableHead = ({
   const columns = useSelector(
     (state: ExperimentsState) => state.tableData.columns
   )
+  const headerDropTargetId = useSelector(
+    (state: ExperimentsState) => state.headerDropTarget
+  )
+  const dispatch = useDispatch()
   const orderedColumns = useColumnOrder(columns, columnOrder)
+  const draggingIds = useRef<string[]>()
+
   const allHeaders: HeaderGroup<Experiment>[] = []
   for (const headerGroup of headerGroups) {
     allHeaders.push(...headerGroup.headers)
   }
 
   const fullColumnOrder = useRef<string[]>()
-  const draggingIds = useRef<string[]>()
   const wrapper = useRef<HTMLDivElement>(null)
 
   useEffect(() => {
@@ -73,26 +80,33 @@ export const TableHead = ({
   }
 
   const onDragUpdate = (e: DragEvent<HTMLElement>) => {
-    const displacer = draggingIds.current
-    displacer &&
-      findDisplacedHeader(e.currentTarget.id, displacedHeader => {
-        const displaced = leafColumnIds(displacedHeader)
-        if (!displaced.some(id => displacer.includes(id))) {
-          fullColumnOrder.current &&
-            setColumnOrder(
-              reorderColumnIds(fullColumnOrder.current, displacer, displaced)
-            )
-        }
-      })
+    findDisplacedHeader(e.currentTarget.id, displacedHeader => {
+      const displaced = leafColumnIds(displacedHeader)
+      dispatch(setDropTarget(displaced[displaced.length - 1]))
+    })
   }
 
   const onDragEnd = () => {
-    draggingIds.current = undefined
     fullColumnOrder.current = undefined
-    sendMessage({
-      payload: columnOrder,
-      type: MessageFromWebviewType.REORDER_COLUMNS
-    })
+    draggingIds.current = undefined
+    dispatch(setDropTarget(''))
+  }
+
+  const onDrop = () => {
+    const fullOrder = fullColumnOrder.current
+    const displacer = draggingIds.current
+    let newOrder: string[] = []
+
+    if (fullOrder && displacer) {
+      newOrder = reorderColumnIds(fullOrder, displacer, [headerDropTargetId])
+
+      setColumnOrder(newOrder)
+      sendMessage({
+        payload: newOrder,
+        type: MessageFromWebviewType.REORDER_COLUMNS
+      })
+      onDragEnd()
+    }
   }
 
   const selectedForPlotsCount = getSelectedForPlotsCount(rows)
@@ -110,6 +124,7 @@ export const TableHead = ({
           onDragStart={onDragStart}
           onDragUpdate={onDragUpdate}
           onDragEnd={onDragEnd}
+          onDrop={onDrop}
           root={root}
           setExpColumnNeedsShadow={setExpColumnNeedsShadow}
         />

webview/src/experiments/components/table/TableHeader.tsx

@@ -31,6 +31,7 @@ interface TableHeaderProps {
   columns: HeaderGroup<Experiment>[]
   orderedColumns: Column[]
   onDragEnter: DragFunction
+  onDragEnd: DragFunction
   onDragStart: DragFunction
   onDrop: DragFunction
   setExpColumnNeedsShadow: (needsShadow: boolean) => void
@@ -42,6 +43,7 @@ export const TableHeader: React.FC<TableHeaderProps> = ({
   columns,
   orderedColumns,
   onDragEnter,
+  onDragEnd,
   onDragStart,
   onDrop,
   root,
@@ -100,6 +102,7 @@ export const TableHeader: React.FC<TableHeaderProps> = ({
       sortEnabled={isSortable}
       hasFilter={hasFilter}
       onDragEnter={onDragEnter}
+      onDragEnd={onDragEnd}
       onDragStart={onDragStart}
       onDrop={onDrop}
       menuDisabled={!isSortable && column.group !== ColumnType.PARAMS}

webview/src/experiments/components/table/TableHeaderCell.tsx

@@ -4,6 +4,7 @@ import {
   ColumnType
 } from 'dvc/src/experiments/webview/contract'
 import React, { useEffect } from 'react'
+import { useSelector } from 'react-redux'
 import { HeaderGroup } from 'react-table'
 import cx from 'classnames'
 import { useInView } from 'react-intersection-observer'
@@ -17,6 +18,7 @@ import {
 } from '../../util/columns'
 import { ContextMenu } from '../../../shared/components/contextMenu/ContextMenu'
 import { DragFunction } from '../../../shared/components/dragDrop/Draggable'
+import { ExperimentsState } from '../../store'
 
 const calcResizerHeight = (
   isPlaceholder: boolean,
@@ -86,6 +88,7 @@ export const TableHeaderCell: React.FC<{
   menuDisabled?: boolean
   menuContent?: React.ReactNode
   onDragEnter: DragFunction
+  onDragEnd: DragFunction
   onDragStart: DragFunction
   onDrop: DragFunction
   setExpColumnNeedsShadow: (needsShadow: boolean) => void
@@ -100,15 +103,20 @@ export const TableHeaderCell: React.FC<{
   menuContent,
   menuDisabled,
   onDragEnter,
+  onDragEnd,
   onDragStart,
   onDrop,
   root,
   setExpColumnNeedsShadow
 }) => {
   const [menuSuppressed, setMenuSuppressed] = React.useState<boolean>(false)
+  const headerDropTargetId = useSelector(
+    (state: ExperimentsState) => state.headerDropTarget
+  )
   const isDraggable = !column.placeholderOf && column.id !== 'id'
 
   const isPlaceholder = !!column.placeholderOf
+
   const canResize = column.canResize && !isPlaceholder
   const resizerHeight = calcResizerHeight(
     isPlaceholder,
@@ -127,6 +135,7 @@ export const TableHeaderCell: React.FC<{
       menuSuppressed={menuSuppressed}
       onDragEnter={onDragEnter}
       onDragStart={onDragStart}
+      onDragEnd={onDragEnd}
       onDrop={onDrop}
       canResize={canResize}
       setMenuSuppressed={setMenuSuppressed}
@@ -159,6 +168,12 @@ export const TableHeaderCell: React.FC<{
         ) : (
           cellContents
         )}
+        <div
+          className={cx(
+            styles.dropTarget,
+            headerDropTargetId === column.id && styles.active
+          )}
+        />
       </div>
     </ContextMenu>
   )

webview/src/experiments/components/table/TableHeaderCellContents.tsx

@@ -34,7 +34,8 @@ export const ColumnDragHandle: React.FC<{
   onDragEnter: DragFunction
   onDragStart: DragFunction
   onDrop: DragFunction
-}> = ({ disabled, column, onDragEnter, onDragStart, onDrop }) => {
+  onDragEnd: DragFunction
+}> = ({ disabled, column, onDragEnter, onDragStart, onDragEnd, onDrop }) => {
   const DropTarget = <span>{column?.name}</span>
 
   return (
@@ -51,6 +52,7 @@ export const ColumnDragHandle: React.FC<{
         dropTarget={DropTarget}
         onDragEnter={onDragEnter}
         onDragStart={onDragStart}
+        onDragEnd={onDragEnd}
         onDrop={onDrop}
       >
         <span>{column?.render('Header')}</span>
@@ -67,6 +69,7 @@ export const TableHeaderCellContents: React.FC<{
   isDraggable: boolean
   menuSuppressed: boolean
   onDragEnter: DragFunction
+  onDragEnd: DragFunction
   onDragStart: DragFunction
   onDrop: DragFunction
   canResize: boolean
@@ -80,6 +83,7 @@ export const TableHeaderCellContents: React.FC<{
   isDraggable,
   menuSuppressed,
   onDragEnter,
+  onDragEnd,
   onDragStart,
   onDrop,
   canResize,
@@ -106,6 +110,7 @@ export const TableHeaderCellContents: React.FC<{
         onDragEnter={onDragEnter}
         onDragStart={onDragStart}
         onDrop={onDrop}
+        onDragEnd={onDragEnd}
       />
       {canResize && (
         <div

webview/src/experiments/components/table/headerDropTargetSlice.ts

@@ -0,0 +1,15 @@
+import { createSlice, PayloadAction } from '@reduxjs/toolkit'
+
+export const headerDropTargetSlice = createSlice({
+  initialState: '',
+  name: 'headerDropTarget',
+  reducers: {
+    setDropTarget: (_, action: PayloadAction<string>) => {
+      return action.payload || ''
+    }
+  }
+})
+
+export const { setDropTarget } = headerDropTargetSlice.actions
+
+export default headerDropTargetSlice.reducer

webview/src/experiments/components/table/styles.module.scss

@@ -509,6 +509,22 @@ $workspace-row-edge-margin: $edge-padding - $cell-padding;
   .th {
     height: auto;
     background-color: $header-bg-color;
+
+    .dropTarget {
+      position: absolute;
+      right: -3px;
+      bottom: 0;
+      width: 3px;
+      z-index: -1;
+      height: var(--table-head-height);
+      background-color: $header-resizer-color;
+      z-index: 1;
+      display: none;
+
+      &.active {
+        display: block;
+      }
+    }
   }
   .td {
     height: auto;

webview/src/experiments/store.ts

@@ -1,10 +1,12 @@
 import { configureStore } from '@reduxjs/toolkit'
 import tableDataReducer from './components/table/tableDataSlice'
 import headersReducer from './components/table/headersSlice'
+import headerDropTargetReducer from './components/table/headerDropTargetSlice'
 import dragAndDropReducer from '../shared/components/dragDrop/dragDropSlice'
 
 export const experimentsReducers = {
   dragAndDrop: dragAndDropReducer,
+  headerDropTarget: headerDropTargetReducer,
   headers: headersReducer,
   tableData: tableDataReducer
 }

webview/src/experiments/util/columns.test.ts

@@ -15,8 +15,8 @@ describe('reorderColumnIds()', () => {
       'id_1'
     ])
     expect(reorderColumnIds(twoColumnIds, ['id_2'], ['id_1'])).toStrictEqual([
-      'id_2',
-      'id_1'
+      'id_1',
+      'id_2'
     ])
 
     const threeColumnIds = [...twoColumnIds, 'id_3']
@@ -41,6 +41,6 @@ describe('reorderColumnIds()', () => {
     ).toStrictEqual(['id_3', 'id_1', 'id_2'])
     expect(
       reorderColumnIds(threeColumnIds, ['id_2', 'id_3'], ['id_1'])
-    ).toStrictEqual(['id_2', 'id_3', 'id_1'])
+    ).toStrictEqual(['id_1', 'id_2', 'id_3'])
   })
 })

webview/src/experiments/util/columns.ts

@@ -96,10 +96,11 @@ export const reorderColumnIds = (
       ...columnIds.slice(displacedIndex + displaced.length)
     ]
   }
+
   return [
     ...columnIds.slice(0, displacedIndex),
-    ...displacer,
     ...displaced,
+    ...displacer,
     ...columnIds.slice(displacedIndex + displaced.length, displacerIndex),
     ...columnIds.slice(displacerIndex + displacer.length)
   ]